"""Walk-forward analysis module for jbqlab.

Walk-forward analysis splits data into in-sample (training) and out-of-sample
(testing) periods to validate strategy robustness.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

import pandas as pd

from jbqlab.engine import run_backtest
from jbqlab.optimize import grid_search

if TYPE_CHECKING:
    from pathlib import Path

__all__ = [
    "WalkForwardResult",
    "run_walk_forward",
    "anchored_walk_forward",
]


@dataclass
class WalkForwardResult:
    """Results from walk-forward analysis."""

    n_folds: int
    in_sample_periods: list[tuple[str, str]]
    out_sample_periods: list[tuple[str, str]]
    best_params_per_fold: list[dict]
    in_sample_metrics: list[dict]
    out_sample_metrics: list[dict]
    combined_out_sample_equity: pd.Series | None
    summary: dict

    def __str__(self) -> str:
        """Human-readable summary."""
        lines = [
            "Walk-Forward Analysis Results",
            "=" * 50,
            f"Number of Folds: {self.n_folds}",
            "",
            "Out-of-Sample Performance:",
            f"  Avg Return: {self.summary['avg_return']:.2%}",
            f"  Avg Sharpe: {self.summary['avg_sharpe']:.2f}",
            f"  Consistency: {self.summary['consistency']:.1%}",
            "",
            "Fold Details:",
        ]

        for i, (params, oos) in enumerate(
            zip(self.best_params_per_fold, self.out_sample_metrics, strict=False)
        ):
            lines.append(f"  Fold {i+1}: params={params}, OOS return={oos.get('total_return', 0):.2%}")

        return "\n".join(lines)


def run_walk_forward(
    data: pd.DataFrame | str | Path,
    strategy: str,
    param_grid: dict[str, list[Any]],
    n_folds: int = 5,
    train_pct: float = 0.8,
    optimize_metric: str = "sharpe",
    min_train_periods: int = 50,
    **kwargs: Any,
) -> WalkForwardResult:
    """Run walk-forward analysis.

    Divides the data into n_folds. For each fold:
    1. Optimize parameters on in-sample (training) data
    2. Test on out-of-sample (testing) data

    Args:
        data: DataFrame or path to CSV.
        strategy: Strategy name.
        param_grid: Parameter grid for optimization.
        n_folds: Number of folds (default 5).
        train_pct: Percentage of each fold for training (default 0.8).
        optimize_metric: Metric to optimize (default 'sharpe').
        min_train_periods: Minimum periods for training (default 50).
        **kwargs: Additional parameters for run_backtest.

    Returns:
        WalkForwardResult with all fold results.
    """
    # Load data if path
    if isinstance(data, (str, Path)):
        from jbqlab.data import load_csv
        df = load_csv(data)
    else:
        df = data.copy()

    n_total = len(df)
    fold_size = n_total // n_folds

    if fold_size < min_train_periods + 10:
        raise ValueError(
            f"Insufficient data for {n_folds} folds. "
            f"Each fold needs at least {min_train_periods + 10} periods."
        )

    in_sample_periods = []
    out_sample_periods = []
    best_params_per_fold = []
    in_sample_metrics = []
    out_sample_metrics = []
    out_sample_equities = []

    for fold in range(n_folds):
        fold_start = fold * fold_size
        fold_end = (fold + 1) * fold_size if fold < n_folds - 1 else n_total

        fold_data = df.iloc[fold_start:fold_end]
        train_end = int(len(fold_data) * train_pct)

        train_data = fold_data.iloc[:train_end]
        test_data = fold_data.iloc[train_end:]

        if len(train_data) < min_train_periods:
            continue

        train_start_date = str(train_data.index[0].date())
        train_end_date = str(train_data.index[-1].date())
        test_start_date = str(test_data.index[0].date())
        test_end_date = str(test_data.index[-1].date())

        in_sample_periods.append((train_start_date, train_end_date))
        out_sample_periods.append((test_start_date, test_end_date))

        # Optimize on training data
        opt_result = grid_search(
            data=train_data,
            strategy=strategy,
            param_grid=param_grid,
            optimize_metric=optimize_metric,
            verbose=False,
            **kwargs,
        )

        best_params = opt_result.best_params
        best_params_per_fold.append(best_params)
        in_sample_metrics.append(opt_result.best_metrics)

        # Test on out-of-sample data
        test_result = run_backtest(
            data=test_data,
            strategy=strategy,
            **best_params,
            **kwargs,
        )

        out_sample_metrics.append(test_result.metrics)
        out_sample_equities.append(test_result.equity_curve)

    # Combine out-of-sample equity curves
    if out_sample_equities:
        combined = pd.concat(out_sample_equities)
        # Normalize to continuous equity
        normalized = combined / combined.iloc[0]
    else:
        normalized = None

    # Calculate summary statistics
    avg_return = sum(m.get("total_return", 0) for m in out_sample_metrics) / len(out_sample_metrics)
    avg_sharpe = sum(m.get("sharpe", 0) for m in out_sample_metrics) / len(out_sample_metrics)
    positive_folds = sum(1 for m in out_sample_metrics if m.get("total_return", 0) > 0)
    consistency = positive_folds / len(out_sample_metrics) if out_sample_metrics else 0

    summary = {
        "avg_return": avg_return,
        "avg_sharpe": avg_sharpe,
        "consistency": consistency,
        "n_positive_folds": positive_folds,
    }

    return WalkForwardResult(
        n_folds=len(out_sample_metrics),
        in_sample_periods=in_sample_periods,
        out_sample_periods=out_sample_periods,
        best_params_per_fold=best_params_per_fold,
        in_sample_metrics=in_sample_metrics,
        out_sample_metrics=out_sample_metrics,
        combined_out_sample_equity=normalized,
        summary=summary,
    )


def anchored_walk_forward(
    data: pd.DataFrame | str | Path,
    strategy: str,
    param_grid: dict[str, list[Any]],
    test_periods: int = 63,  # ~3 months
    reoptimize_every: int = 21,  # ~1 month
    min_train_periods: int = 252,  # ~1 year
    optimize_metric: str = "sharpe",
    **kwargs: Any,
) -> WalkForwardResult:
    """Run anchored walk-forward analysis.

    Uses expanding window for training (anchored to start).
    Re-optimizes parameters at regular intervals.

    Args:
        data: DataFrame or path to CSV.
        strategy: Strategy name.
        param_grid: Parameter grid.
        test_periods: Out-of-sample test length per iteration.
        reoptimize_every: Re-optimize after this many periods.
        min_train_periods: Minimum training periods before first test.
        optimize_metric: Metric to optimize.
        **kwargs: Additional parameters.

    Returns:
        WalkForwardResult with anchored analysis.
    """
    # Load data if path
    if isinstance(data, (str, Path)):
        from jbqlab.data import load_csv
        df = load_csv(data)
    else:
        df = data.copy()

    n_total = len(df)

    if n_total < min_train_periods + test_periods:
        raise ValueError("Insufficient data for anchored walk-forward")

    in_sample_periods = []
    out_sample_periods = []
    best_params_per_fold = []
    in_sample_metrics = []
    out_sample_metrics = []
    out_sample_equities = []

    train_end = min_train_periods
    current_params = None
    periods_since_optimize = 0

    while train_end < n_total - test_periods:
        # Check if we need to re-optimize
        if current_params is None or periods_since_optimize >= reoptimize_every:
            train_data = df.iloc[:train_end]

            opt_result = grid_search(
                data=train_data,
                strategy=strategy,
                param_grid=param_grid,
                optimize_metric=optimize_metric,
                verbose=False,
                **kwargs,
            )

            current_params = opt_result.best_params
            best_params_per_fold.append(current_params)
            in_sample_metrics.append(opt_result.best_metrics)

            train_start_date = str(train_data.index[0].date())
            train_end_date = str(train_data.index[-1].date())
            in_sample_periods.append((train_start_date, train_end_date))

            periods_since_optimize = 0

        # Test on next period
        test_start = train_end
        test_end = min(train_end + test_periods, n_total)
        test_data = df.iloc[test_start:test_end]

        test_result = run_backtest(
            data=test_data,
            strategy=strategy,
            **current_params,
            **kwargs,
        )

        test_start_date = str(test_data.index[0].date())
        test_end_date = str(test_data.index[-1].date())
        out_sample_periods.append((test_start_date, test_end_date))
        out_sample_metrics.append(test_result.metrics)
        out_sample_equities.append(test_result.equity_curve)

        # Move forward
        train_end = test_end
        periods_since_optimize += test_periods

    # Combine out-of-sample equity
    if out_sample_equities:
        combined = pd.concat(out_sample_equities)
        normalized = combined / combined.iloc[0]
    else:
        normalized = None

    # Summary
    avg_return = sum(m.get("total_return", 0) for m in out_sample_metrics) / max(len(out_sample_metrics), 1)
    avg_sharpe = sum(m.get("sharpe", 0) for m in out_sample_metrics) / max(len(out_sample_metrics), 1)
    positive_folds = sum(1 for m in out_sample_metrics if m.get("total_return", 0) > 0)
    consistency = positive_folds / max(len(out_sample_metrics), 1)

    summary = {
        "avg_return": avg_return,
        "avg_sharpe": avg_sharpe,
        "consistency": consistency,
        "n_positive_folds": positive_folds,
    }

    return WalkForwardResult(
        n_folds=len(out_sample_metrics),
        in_sample_periods=in_sample_periods,
        out_sample_periods=out_sample_periods,
        best_params_per_fold=best_params_per_fold,
        in_sample_metrics=in_sample_metrics,
        out_sample_metrics=out_sample_metrics,
        combined_out_sample_equity=normalized,
        summary=summary,
    )
