"""Walk-forward validation for Vanta Crunch signal models.

Answers the question: **is the signal real, or overfit to the backtest period?**

Splits data into sequential folds and runs the model independently on each,
evaluating only on the test window.  Aggregates per-fold metrics to measure
signal stability.

Architecture
────────────
This module orchestrates ``BacktestRunner`` — it doesn't replicate it.

    validation.py   ←  "Is the signal consistent across time?"
        ↓ uses
    backtest.py     ←  "What's the signal over one period?"
        ↓ feeds
    miner/backtest  ←  "How does the position layer perform?"

Key design decisions:

- **model_factory** (callable) instead of model instance — models are stateful
  (indicators accumulate), so each fold needs a fresh instance.
- **Warmup period** before each test window — indicators need history to
  bootstrap.  Warmup data is fed via tick() but predictions during warmup
  are discarded.
- **Per-fold isolation** — no state carries between folds.  This is stricter
  than production (where indicators run continuously) but reveals whether
  the signal depends on specific history.

Usage::

    from vanta.validation import WalkForwardValidator, expanding_splits, rolling_splits
    from vanta.examples import HorizonAwareTracker

    validator = WalkForwardValidator(
        model_factory=HorizonAwareTracker,
        prediction_interval_seconds=43200,
        horizon_seconds=86400,
    )

    # 5 expanding folds (train grows, test stays ~2 months)
    folds = expanding_splits("2025-02-15", "2026-01-01", n_test_days=60, n_folds=5)

    result = validator.run(folds)
    result.summary()     # per-fold IC, aggregate stability
    result.plot()        # IC over time, fold comparison
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from typing import Any, Callable

from vanta.backtest import BacktestResult, BacktestRunner


# ── Fold definitions ──────────────────────────────────────────────────────────


@dataclass(frozen=True)
class Fold:
    """One train/test split for walk-forward validation.

    The model runs on data from ``warmup_start`` to ``test_end``.
    Only predictions within ``test_start`` to ``test_end`` are scored.
    The warmup period (warmup_start → test_start) bootstraps indicators.
    """

    warmup_start: str  # ISO date — beginning of data fed to model
    test_start: str  # ISO date — beginning of scored evaluation
    test_end: str  # ISO date — end of scored evaluation
    name: str = ""  # human-readable label

    @property
    def warmup_days(self) -> int:
        ws = datetime.strptime(self.warmup_start, "%Y-%m-%d")
        ts = datetime.strptime(self.test_start, "%Y-%m-%d")
        return (ts - ws).days

    @property
    def test_days(self) -> int:
        ts = datetime.strptime(self.test_start, "%Y-%m-%d")
        te = datetime.strptime(self.test_end, "%Y-%m-%d")
        return (te - ts).days


def rolling_splits(
    data_start: str,
    data_end: str,
    warmup_days: int = 30,
    test_days: int = 60,
    step_days: int | None = None,
) -> list[Fold]:
    """Generate rolling-window folds.

    Each fold has a fixed-size warmup and test window that slides forward.

    Args:
        data_start: First available data date (ISO).
        data_end: Last available data date (ISO).
        warmup_days: Days of warmup before each test window.
        test_days: Days in each test window.
        step_days: How far to slide between folds. Defaults to ``test_days``
            (non-overlapping test windows).

    Returns:
        List of Fold objects.
    """
    if step_days is None:
        step_days = test_days

    start = datetime.strptime(data_start, "%Y-%m-%d")
    end = datetime.strptime(data_end, "%Y-%m-%d")

    folds = []
    fold_num = 1
    warmup_start = start

    while True:
        test_start = warmup_start + timedelta(days=warmup_days)
        test_end = test_start + timedelta(days=test_days)

        if test_end > end:
            break

        folds.append(Fold(
            warmup_start=warmup_start.strftime("%Y-%m-%d"),
            test_start=test_start.strftime("%Y-%m-%d"),
            test_end=test_end.strftime("%Y-%m-%d"),
            name=f"fold_{fold_num}",
        ))
        fold_num += 1
        warmup_start += timedelta(days=step_days)

    return folds


def expanding_splits(
    data_start: str,
    data_end: str,
    n_test_days: int = 60,
    n_folds: int = 5,
    min_warmup_days: int = 30,
) -> list[Fold]:
    """Generate expanding-window folds.

    Warmup grows with each fold (more history), test window is fixed size.
    The last fold uses all available data minus the test window.

    Args:
        data_start: First available data date (ISO).
        data_end: Last available data date (ISO).
        n_test_days: Days in each test window.
        n_folds: Number of folds to generate.
        min_warmup_days: Minimum warmup for the first fold.

    Returns:
        List of Fold objects.
    """
    start = datetime.strptime(data_start, "%Y-%m-%d")
    end = datetime.strptime(data_end, "%Y-%m-%d")
    total_days = (end - start).days

    # Work backwards from the end: last fold's test_end = data_end
    # Space test windows evenly across available data after min_warmup
    available_for_tests = total_days - min_warmup_days
    if available_for_tests < n_test_days:
        raise ValueError(
            f"Not enough data: {total_days} days available, need at least "
            f"{min_warmup_days + n_test_days} for 1 fold"
        )

    # Step between test_start dates
    step = (available_for_tests - n_test_days) / max(n_folds - 1, 1)

    folds = []
    for i in range(n_folds):
        test_start = start + timedelta(days=min_warmup_days + int(i * step))
        test_end = test_start + timedelta(days=n_test_days)

        if test_end > end:
            test_end = end
        if test_start >= test_end:
            break

        folds.append(Fold(
            warmup_start=data_start,  # expanding: always start from beginning
            test_start=test_start.strftime("%Y-%m-%d"),
            test_end=test_end.strftime("%Y-%m-%d"),
            name=f"fold_{i + 1}",
        ))

    return folds


# ── Validator ─────────────────────────────────────────────────────────────────


class WalkForwardValidator:
    """Run a model across walk-forward folds, aggregate stability metrics.

    Args:
        model_factory: Callable that returns a fresh model instance.
            Called once per fold.  Can be a class (``HorizonAwareTracker``)
            or a lambda (``lambda: HorizonAwareTracker(atr_period=10)``).
        prediction_interval_seconds: How often to predict (passed to BacktestRunner).
        horizon_seconds: Prediction horizon (passed to BacktestRunner).
        window_size: Candle warmup window for BacktestRunner.
        transaction_cost_bps: Cost in bps for turnover.
        **runner_kwargs: Extra kwargs passed to ``BacktestRunner.run()``.
    """

    def __init__(
        self,
        model_factory: Callable[[], Any],
        prediction_interval_seconds: int = 43200,
        horizon_seconds: int = 86400,
        window_size: int = 120,
        transaction_cost_bps: float = 0.0,
        **runner_kwargs: Any,
    ):
        self.model_factory = model_factory
        self.prediction_interval_seconds = prediction_interval_seconds
        self.horizon_seconds = horizon_seconds
        self.window_size = window_size
        self.transaction_cost_bps = transaction_cost_bps
        self.runner_kwargs = runner_kwargs

    def run(
        self,
        folds: list[Fold],
        subject: str = "BTCUSDT",
        verbose: bool = True,
    ) -> WalkForwardResult:
        """Execute walk-forward validation across all folds.

        Each fold gets a fresh model instance. The model runs from
        ``fold.warmup_start`` to ``fold.test_end``, but only predictions
        within the test window are included in that fold's metrics.

        Args:
            folds: List of Fold definitions.
            subject: Trading pair.
            verbose: Print progress per fold.

        Returns:
            WalkForwardResult with per-fold and aggregate metrics.
        """
        fold_results: list[FoldResult] = []

        for i, fold in enumerate(folds):
            if verbose:
                print(
                    f"  Fold {i + 1}/{len(folds)}: "
                    f"{fold.test_start} → {fold.test_end} "
                    f"(warmup={fold.warmup_days}d, test={fold.test_days}d)...",
                    end="",
                    flush=True,
                )

            # Fresh model per fold — no state leakage
            model = self.model_factory()

            runner = BacktestRunner(
                model=model,
                transaction_cost_bps=self.transaction_cost_bps,
            )

            # Run from warmup_start → test_end (model sees warmup data,
            # predictions scored from warmup_start but we'll filter)
            bt_result = runner.run(
                subject=subject,
                start=fold.warmup_start,
                end=fold.test_end,
                window_size=self.window_size,
                prediction_interval_seconds=self.prediction_interval_seconds,
                horizon_seconds=self.horizon_seconds,
                **self.runner_kwargs,
            )

            # Filter predictions to test window only
            test_start_dt = datetime.strptime(fold.test_start, "%Y-%m-%d").replace(
                tzinfo=timezone.utc
            )
            test_end_dt = datetime.strptime(fold.test_end, "%Y-%m-%d").replace(
                tzinfo=timezone.utc
            )

            test_preds = []
            for p in bt_result._predictions:
                ts = p.get("ts")
                if ts is None:
                    continue
                if not hasattr(ts, "timestamp"):
                    ts = datetime.fromtimestamp(ts, tz=timezone.utc)
                if test_start_dt <= ts <= test_end_dt:
                    test_preds.append(p)

            # Compute fold-level metrics from test predictions only
            fold_metrics = _compute_fold_metrics(test_preds)

            fold_result = FoldResult(
                fold=fold,
                backtest_result=bt_result,
                test_predictions=test_preds,
                metrics=fold_metrics,
            )
            fold_results.append(fold_result)

            if verbose:
                ic = fold_metrics.get("ic", 0)
                n = fold_metrics.get("n_predictions", 0)
                hr = fold_metrics.get("hit_rate", 0)
                print(f"  IC={ic:+.4f}  HR={hr:.1%}  N={n}")

        return WalkForwardResult(fold_results=fold_results)


# ── Results ───────────────────────────────────────────────────────────────────


@dataclass
class FoldResult:
    """Results for one fold of walk-forward validation."""

    fold: Fold
    backtest_result: BacktestResult
    test_predictions: list[dict]
    metrics: dict[str, float]


@dataclass
class WalkForwardResult:
    """Aggregate walk-forward validation results."""

    fold_results: list[FoldResult]

    @property
    def n_folds(self) -> int:
        return len(self.fold_results)

    @property
    def fold_metrics(self) -> list[dict[str, float]]:
        """Per-fold metrics as a list."""
        return [fr.metrics for fr in self.fold_results]

    @property
    def aggregate(self) -> dict[str, float]:
        """Aggregate statistics across folds."""
        if not self.fold_results:
            return {}

        metrics_keys = ["ic", "hit_rate", "mean_score", "sharpe", "n_predictions"]
        agg: dict[str, float] = {}

        for key in metrics_keys:
            values = [fm.get(key, 0.0) for fm in self.fold_metrics]
            n = len(values)
            mean = sum(values) / n
            agg[f"{key}_mean"] = mean
            if n > 1:
                var = sum((v - mean) ** 2 for v in values) / (n - 1)
                agg[f"{key}_std"] = math.sqrt(var)
            else:
                agg[f"{key}_std"] = 0.0
            agg[f"{key}_min"] = min(values)
            agg[f"{key}_max"] = max(values)

        # Stability metrics
        ics = [fm.get("ic", 0.0) for fm in self.fold_metrics]
        ic_mean = agg.get("ic_mean", 0.0)
        ic_std = agg.get("ic_std", 1e-9)

        # IC information ratio: mean / std — how consistent is the signal?
        agg["ic_ir"] = ic_mean / ic_std if ic_std > 1e-9 else 0.0

        # % of folds with positive IC
        agg["ic_positive_pct"] = sum(1 for ic in ics if ic > 0) / len(ics) if ics else 0.0

        # IC decay: correlation of IC with fold index (time)
        # Negative = signal degrades over time (possible overfit to early data)
        if len(ics) >= 3:
            indices = list(range(len(ics)))
            agg["ic_time_corr"] = _correlation(
                [float(x) for x in indices], ics
            )
        else:
            agg["ic_time_corr"] = 0.0

        # Total predictions across all folds
        agg["total_predictions"] = sum(
            fm.get("n_predictions", 0) for fm in self.fold_metrics
        )

        return agg

    def summary(self) -> str:
        """Print formatted walk-forward summary."""
        lines = [
            "═" * 70,
            "  WALK-FORWARD VALIDATION",
            "═" * 70,
            f"  Folds: {self.n_folds}",
        ]

        # Per-fold table
        lines.append("─" * 70)
        header = f"  {'Fold':<12} {'Period':<25} {'IC':>8} {'HR':>8} {'Sharpe':>8} {'N':>6}"
        lines.append(header)
        lines.append("─" * 70)

        for fr in self.fold_results:
            f = fr.fold
            m = fr.metrics
            period = f"{f.test_start} → {f.test_end}"
            lines.append(
                f"  {f.name:<12} {period:<25} "
                f"{m.get('ic', 0):>+8.4f} "
                f"{m.get('hit_rate', 0):>7.1%} "
                f"{m.get('sharpe', 0):>+8.2f} "
                f"{m.get('n_predictions', 0):>6.0f}"
            )

        # Aggregate
        agg = self.aggregate
        lines.append("─" * 70)
        lines.append("  AGGREGATE:")
        lines.append(
            f"    IC:       {agg.get('ic_mean', 0):+.4f} ± {agg.get('ic_std', 0):.4f}"
            f"   (min={agg.get('ic_min', 0):+.4f}, max={agg.get('ic_max', 0):+.4f})"
        )
        lines.append(
            f"    Hit Rate: {agg.get('hit_rate_mean', 0):.1%} ± {agg.get('hit_rate_std', 0):.1%}"
        )
        lines.append(
            f"    Sharpe:   {agg.get('sharpe_mean', 0):+.2f} ± {agg.get('sharpe_std', 0):.2f}"
        )

        # Stability verdict
        lines.append("─" * 70)
        lines.append("  STABILITY:")
        ir = agg.get("ic_ir", 0)
        pos_pct = agg.get("ic_positive_pct", 0)
        time_corr = agg.get("ic_time_corr", 0)

        ic_ir_verdict = "✅ stable" if ir > 0.5 else "⚠️  marginal" if ir > 0.2 else "❌ unstable"
        pos_verdict = "✅" if pos_pct >= 0.8 else "⚠️ " if pos_pct >= 0.6 else "❌"
        decay_verdict = "✅ no decay" if time_corr > -0.5 else "⚠️  some decay" if time_corr > -0.8 else "❌ decaying"

        lines.append(f"    IC IR (mean/std):     {ir:+.2f}  {ic_ir_verdict}")
        lines.append(f"    IC positive folds:    {pos_pct:.0%}  {pos_verdict}")
        lines.append(f"    IC time correlation:  {time_corr:+.2f}  {decay_verdict}")

        overall = "✅ SIGNAL APPEARS ROBUST" if (ir > 0.5 and pos_pct >= 0.8) else \
                  "⚠️  SIGNAL NEEDS MORE DATA" if (ir > 0.2 or pos_pct >= 0.6) else \
                  "❌ SIGNAL MAY BE OVERFIT"
        lines.append("")
        lines.append(f"  VERDICT: {overall}")
        lines.append("═" * 70)

        text = "\n".join(lines)
        print(text)
        return text

    def plot(self, save: str | None = None):
        """Plot per-fold IC and cumulative metrics."""
        try:
            import matplotlib.pyplot as plt
        except ImportError:
            print("pip install matplotlib")
            return

        fig, axes = plt.subplots(2, 2, figsize=(14, 10))
        fig.suptitle("Walk-Forward Validation", fontsize=14)

        fold_names = [fr.fold.name for fr in self.fold_results]
        ics = [fr.metrics.get("ic", 0) for fr in self.fold_results]
        hrs = [fr.metrics.get("hit_rate", 0) for fr in self.fold_results]
        sharpes = [fr.metrics.get("sharpe", 0) for fr in self.fold_results]
        ns = [fr.metrics.get("n_predictions", 0) for fr in self.fold_results]

        # IC per fold
        ax = axes[0, 0]
        colors = ["green" if ic > 0 else "red" for ic in ics]
        ax.bar(fold_names, ics, color=colors, alpha=0.7)
        ax.axhline(y=0, color="gray", linestyle="--", linewidth=0.5)
        ax.set_title("IC per Fold")
        ax.set_ylabel("Information Coefficient")
        ax.tick_params(axis="x", rotation=45)

        # Hit rate per fold
        ax = axes[0, 1]
        ax.bar(fold_names, hrs, color="steelblue", alpha=0.7)
        ax.axhline(y=0.5, color="gray", linestyle="--", linewidth=0.5)
        ax.set_title("Hit Rate per Fold")
        ax.set_ylabel("Hit Rate")
        ax.set_ylim(0, 1)
        ax.tick_params(axis="x", rotation=45)

        # Cumulative score per fold (equity curves overlaid)
        ax = axes[1, 0]
        for fr in self.fold_results:
            scored = [p for p in fr.test_predictions if p.get("score") is not None]
            if not scored:
                continue
            cum = []
            running = 0.0
            for p in scored:
                running += p["score"]
                cum.append(running)
            ax.plot(range(len(cum)), cum, label=fr.fold.name, linewidth=1)
        ax.axhline(y=0, color="gray", linestyle="--", linewidth=0.5)
        ax.set_title("Cumulative Score per Fold")
        ax.set_xlabel("Prediction #")
        ax.set_ylabel("Cumulative Score")
        ax.legend(fontsize=7)

        # Sharpe per fold
        ax = axes[1, 1]
        colors = ["green" if s > 0 else "red" for s in sharpes]
        ax.bar(fold_names, sharpes, color=colors, alpha=0.7)
        ax.axhline(y=0, color="gray", linestyle="--", linewidth=0.5)
        ax.set_title("Sharpe per Fold")
        ax.set_ylabel("Sharpe Ratio")
        ax.tick_params(axis="x", rotation=45)

        plt.tight_layout()
        if save:
            plt.savefig(save, dpi=150, bbox_inches="tight")
            print(f"Saved to {save}")
        else:
            plt.show()

    def to_dataframe(self):
        """Convert fold metrics to a pandas DataFrame."""
        import pandas as pd
        rows = []
        for fr in self.fold_results:
            row = {
                "fold": fr.fold.name,
                "test_start": fr.fold.test_start,
                "test_end": fr.fold.test_end,
                "warmup_days": fr.fold.warmup_days,
                "test_days": fr.fold.test_days,
                **fr.metrics,
            }
            rows.append(row)
        return pd.DataFrame(rows)


# ── Fold-level metric computation ─────────────────────────────────────────────


def _compute_fold_metrics(predictions: list[dict]) -> dict[str, float]:
    """Compute signal-quality metrics from a list of predictions.

    Only uses scored predictions (those with actual_return).
    """
    scored = [p for p in predictions if p.get("actual_return") is not None]

    metrics: dict[str, float] = {
        "n_predictions": len(scored),
    }

    if not scored:
        metrics.update(ic=0.0, hit_rate=0.0, mean_score=0.0, sharpe=0.0)
        return metrics

    expected = [p.get("expected_return", 0.0) for p in scored]
    actual = [p.get("actual_return", 0.0) for p in scored]
    scores = [p.get("score", 0.0) for p in scored]

    # IC
    metrics["ic"] = _correlation(expected, actual) if len(scored) >= 3 else 0.0

    # Hit rate
    directional = [p for p in scored if p.get("direction_correct") is not None]
    if directional:
        metrics["hit_rate"] = sum(1 for p in directional if p["direction_correct"]) / len(directional)
    else:
        metrics["hit_rate"] = 0.0

    # Mean score
    metrics["mean_score"] = sum(scores) / len(scores)

    # Sharpe of scores
    n = len(scores)
    mean = metrics["mean_score"]
    if n > 1:
        var = sum((s - mean) ** 2 for s in scores) / (n - 1)
        std = math.sqrt(var)
        # Annualise assuming scores are per prediction_interval
        metrics["sharpe"] = (mean / std * math.sqrt(n)) if std > 1e-12 else 0.0
    else:
        metrics["sharpe"] = 0.0

    # Score stats
    metrics["total_score"] = sum(scores)
    metrics["score_std"] = math.sqrt(sum((s - mean) ** 2 for s in scores) / max(n - 1, 1)) if n > 1 else 0.0

    return metrics


def _correlation(x: list[float], y: list[float]) -> float:
    """Pearson correlation."""
    n = len(x)
    if n < 2:
        return 0.0
    mean_x = sum(x) / n
    mean_y = sum(y) / n
    cov = sum((xi - mean_x) * (yi - mean_y) for xi, yi in zip(x, y)) / n
    std_x = (sum((xi - mean_x) ** 2 for xi in x) / n) ** 0.5
    std_y = (sum((yi - mean_y) ** 2 for yi in y) / n) ** 0.5
    if std_x < 1e-12 or std_y < 1e-12:
        return 0.0
    return cov / (std_x * std_y)
