"""Walk-forward / CPCV splitting with purge and embargo.

Anchored on López de Prado, Advances in Financial Machine Learning (2018),
ch. 7 and ch. 11-12, implemented on top of the `purgedcv` package (MIT,
pypi purgedcv==0.1.11).

Conventions used by this eval:
  - labels carry a horizon h (forward return over h bars); training samples
    whose label window [t, t+h] overlaps a test block are PURGED
  - EMBARGO of e bars is additionally applied after each test block;
    default e = max(label horizon, one funding interval in bars) = 8 bars at 1h
  - CPCV: N blocks, K held out per combination -> C(N, K) test sets, which
    recombine into backtest paths via purgedcv.reconstruct_paths.
    Defaults N=6, K=2: 15 combinations, C(5,1)=5 paths.
"""

from __future__ import annotations

import numpy as np
import pandas as pd

try:
    from purgedcv import (
        CombinatorialPurgedCV,
        reconstruct_paths,
    )
except ImportError as e:  # pragma: no cover
    raise ImportError(
        "purgedcv is required: pip install purgedcv==0.1.11"
    ) from e

__all__ = ["make_cpcv", "make_walkforward", "SplitAudit"]


class SplitAudit:
    """No-lookahead audit for a produced split."""

    def __init__(self, n_samples: int, horizon: int, embargo: int):
        self.n_samples = n_samples
        self.horizon = horizon
        self.embargo = embargo
        self.violations: list[str] = []

    def check(self, train_idx: np.ndarray, test_idx: np.ndarray, label: str) -> None:
        if len(test_idx) == 0 or len(train_idx) == 0:
            self.violations.append(f"{label}: empty fold")
            return
        # CPCV test sets are the UNION of K blocks: train samples between
        # two test blocks are legitimate. Audit each contiguous test
        # segment separately instead of the whole [min, max] range.
        breaks = np.flatnonzero(np.diff(np.sort(test_idx)) > 1)
        segments = np.split(np.sort(test_idx), breaks + 1)
        total_bad = 0
        for seg in segments:
            t_min, t_max = int(seg.min()), int(seg.max())
            # overlap iff i <= t_max+embargo and i+h >= t_min
            leak_zone_end = t_max + self.embargo
            bad = int(((train_idx <= leak_zone_end)
                       & (train_idx + self.horizon >= t_min)).sum())
            total_bad += bad
        if total_bad:
            self.violations.append(
                f"{label}: {total_bad} training samples overlap test window "
                f"(leakage)"
            )

    @property
    def ok(self) -> bool:
        return not self.violations


def make_cpcv(index: pd.DatetimeIndex, horizon: int, embargo: int,
              freq: str = "1h", n_blocks: int = 6, k: int = 2):
    """Return the purgedcv CPCV splitter and audit it.

    index: bar start times (UTC). horizon/embargo are in BARS.
    purgedcv 0.1.11 requires prediction/evaluation times and takes
    purge/embargo as timedeltas; it purges internally, SplitAudit is
    kept as an independent leakage check.
    """
    freq_td = pd.Timedelta(freq)
    horizon_td = horizon * freq_td
    embargo_td = embargo * freq_td
    if embargo_td < horizon_td:
        raise ValueError(
            f"embargo ({embargo} bars) must be >= label horizon ({horizon} bars)"
        )
    splitter = CombinatorialPurgedCV(
        n_splits=n_blocks,
        n_test_groups=k,
        prediction_times=index.values,
        evaluation_times=(index + freq_td).values,
        purge_horizon=horizon_td,
        embargo=embargo_td,
    )
    n_samples = len(index)
    audit = SplitAudit(n_samples, horizon, embargo)
    combos = list(splitter.split(np.arange(n_samples)))
    for j, (train_idx, test_idx) in enumerate(combos):
        audit.check(np.asarray(train_idx), np.asarray(test_idx), f"cpcv-combo-{j}")
    if not audit.ok:
        raise RuntimeError(
            "no-lookahead audit failed:\n" + "\n".join(audit.violations)
        )
    return splitter, combos, audit


def make_walkforward(n_samples: int, n_folds: int, horizon: int, embargo: int):
    """Expanding-window walk-forward with purge + embargo, audited."""
    if embargo < horizon:
        raise ValueError(
            f"embargo ({embargo}) must be >= label horizon ({horizon})"
        )
    fold_bounds = np.linspace(0, n_samples, n_folds + 2, dtype=int)
    folds = []
    audit = SplitAudit(n_samples, horizon, embargo)
    for j in range(n_folds):
        test_start, test_end = fold_bounds[j + 1], fold_bounds[j + 2]
        train_idx = np.arange(0, test_start - embargo - horizon)
        test_idx = np.arange(test_start, test_end)
        audit.check(train_idx, test_idx, f"wf-fold-{j}")
        folds.append((train_idx, test_idx))
    if not audit.ok:
        raise RuntimeError(
            "no-lookahead audit failed:\n" + "\n".join(audit.violations)
        )
    return folds, audit