"""Walk-forward optimization for backtesting."""

from __future__ import annotations

import itertools
from dataclasses import dataclass, field
from typing import Any, Callable

from horizon.analytics import Metrics, compute_metrics
from horizon.backtest import backtest, Tick
from horizon.result import BacktestResult


@dataclass
class WalkForwardWindow:
    """A single train/test window in walk-forward optimization."""

    train_start: float
    train_end: float
    test_start: float
    test_end: float


@dataclass
class WalkForwardResult:
    """Result of walk-forward optimization."""

    windows: list[WalkForwardWindow] = field(default_factory=list)
    best_params_per_window: list[dict[str, Any]] = field(default_factory=list)
    test_results: list[BacktestResult] = field(default_factory=list)
    aggregate_equity: list[tuple[float, float]] = field(default_factory=list)
    aggregate_metrics: Metrics = field(default_factory=Metrics)


def walk_forward(
    data: Any,
    pipeline_factory: Callable[[dict], list[Callable]],
    param_grid: dict[str, list[Any]],
    n_splits: int = 5,
    train_ratio: float = 0.7,
    expanding: bool = True,
    objective: str = "sharpe_ratio",
    purge_gap: float = 0.0,
    markets: list[str] | None = None,
    risk: Any = None,
    initial_capital: float = 1000.0,
    **backtest_kwargs: Any,
) -> WalkForwardResult:
    """Run walk-forward optimization over historical data.

    Args:
        data: Historical data (same formats as hz.backtest).
        pipeline_factory: Function that takes params dict and returns a pipeline.
        param_grid: Dict of parameter names to lists of values for grid search.
        n_splits: Number of train/test splits.
        train_ratio: Fraction of each split used for training (rest is test).
        expanding: If True, training window expands from the start (anchored).
                   If False, training window rolls forward (fixed width).
        objective: Metric name to optimize (attribute on Metrics).
        purge_gap: Seconds to purge between train/test to prevent lookahead.
        markets: Market IDs (passed to backtest).
        risk: Risk config (passed to backtest).
        initial_capital: Starting capital.
        **backtest_kwargs: Additional kwargs passed to each backtest() call.

    Returns:
        WalkForwardResult with per-window results and aggregate metrics.
    """
    from horizon.backtest import _normalize_data

    # Normalize data to ticks for timestamp access
    if isinstance(data, dict):
        # Multi-feed: use first feed for timestamp extraction
        first_key = next(iter(data))
        ticks = _normalize_data(data[first_key])
    else:
        ticks = _normalize_data(data)

    if not ticks:
        return WalkForwardResult()

    timestamps = [t.timestamp for t in ticks]
    ts_start = timestamps[0]
    ts_end = timestamps[-1]

    # Generate windows
    windows = _generate_windows(ts_start, ts_end, n_splits, train_ratio, expanding, purge_gap)
    if not windows:
        return WalkForwardResult()

    # Generate all parameter combinations
    param_names = list(param_grid.keys())
    param_values = list(param_grid.values())
    param_combos = [dict(zip(param_names, combo)) for combo in itertools.product(*param_values)]

    if not param_combos:
        return WalkForwardResult()

    result = WalkForwardResult(windows=windows)

    for window in windows:
        # Grid search on training data
        best_score = None
        best_params = param_combos[0]

        for params in param_combos:
            pipeline = pipeline_factory(params)
            train_data = _filter_data_by_time(data, window.train_start, window.train_end)

            train_result = backtest(
                data=train_data,
                pipeline=pipeline,
                markets=markets,
                risk=risk,
                initial_capital=initial_capital,
                **backtest_kwargs,
            )

            score = getattr(train_result.metrics, objective, 0.0)
            if score is None:
                score = 0.0

            if best_score is None or score > best_score:
                best_score = score
                best_params = params

        result.best_params_per_window.append(best_params)

        # Run best params on test window
        test_pipeline = pipeline_factory(best_params)
        test_data = _filter_data_by_time(data, window.test_start, window.test_end)

        test_result = backtest(
            data=test_data,
            pipeline=test_pipeline,
            markets=markets,
            risk=risk,
            initial_capital=initial_capital,
            **backtest_kwargs,
        )
        result.test_results.append(test_result)

    # Aggregate: chain test equity curves
    aggregate_equity: list[tuple[float, float]] = []
    running_capital = initial_capital

    for test_result in result.test_results:
        if not test_result.equity_curve:
            continue
        # Offset equity to chain from previous ending capital
        first_eq = test_result.equity_curve[0][1]
        offset = running_capital - first_eq
        for ts, eq in test_result.equity_curve:
            aggregate_equity.append((ts, eq + offset))
        running_capital = aggregate_equity[-1][1] if aggregate_equity else running_capital

    result.aggregate_equity = aggregate_equity

    # Compute aggregate metrics
    all_trades = []
    for tr in result.test_results:
        all_trades.extend(tr.trades)

    result.aggregate_metrics = compute_metrics(
        aggregate_equity, all_trades, initial_capital
    )

    return result


def _generate_windows(
    ts_start: float,
    ts_end: float,
    n_splits: int,
    train_ratio: float,
    expanding: bool,
    purge_gap: float,
) -> list[WalkForwardWindow]:
    """Generate train/test windows for walk-forward optimization."""
    total_duration = ts_end - ts_start
    if total_duration <= 0 or n_splits <= 0:
        return []

    # Each split covers total_duration / n_splits
    split_duration = total_duration / n_splits
    windows = []

    for i in range(n_splits):
        split_end = ts_start + (i + 1) * split_duration

        if expanding:
            # Anchored expanding: training always starts from ts_start
            train_start = ts_start
        else:
            # Rolling: training window has fixed width
            train_start = ts_start + i * split_duration * (1 - train_ratio)

        # Train/test boundary within this split
        train_end = ts_start + i * split_duration + split_duration * train_ratio
        test_start = train_end + purge_gap
        test_end = split_end

        if test_start >= test_end:
            continue

        windows.append(WalkForwardWindow(
            train_start=train_start,
            train_end=train_end,
            test_start=test_start,
            test_end=test_end,
        ))

    return windows


def _filter_data_by_time(data: Any, start: float, end: float) -> Any:
    """Filter data to a time range.

    Handles list[dict], dict[str, list[dict]], and list[Tick].
    """
    if isinstance(data, dict):
        # Multi-feed: filter each feed
        return {
            key: _filter_data_by_time(val, start, end) for key, val in data.items()
        }

    if isinstance(data, list):
        if not data:
            return data
        # Check if it's list of dicts or list of Ticks
        first = data[0]
        if isinstance(first, dict):
            return [
                rec for rec in data
                if start <= float(rec.get("timestamp", 0)) <= end
            ]
        elif isinstance(first, Tick):
            return [t for t in data if start <= t.timestamp <= end]
        else:
            # Try duck-typing with safe attribute/key access
            return [
                rec for rec in data
                if start <= float(
                    getattr(rec, "timestamp", None)
                    if hasattr(rec, "timestamp")
                    else rec.get("timestamp", 0) if hasattr(rec, "get") else 0
                ) <= end
            ]

    return data
