"""Tests for walk-forward optimization."""

import pytest
from horizon._horizon import Quote
from horizon.walkforward import (
    walk_forward,
    _generate_windows,
    _filter_data_by_time,
    WalkForwardWindow,
    WalkForwardResult,
)
from horizon.backtest import Tick


class TestGenerateWindows:
    def test_basic_windows(self):
        windows = _generate_windows(0.0, 100.0, 5, 0.7, expanding=True, purge_gap=0.0)
        assert len(windows) == 5
        for w in windows:
            assert w.train_start <= w.train_end
            assert w.test_start <= w.test_end

    def test_expanding_windows_anchor_at_start(self):
        windows = _generate_windows(0.0, 100.0, 3, 0.7, expanding=True, purge_gap=0.0)
        for w in windows:
            assert w.train_start == 0.0

    def test_rolling_windows(self):
        windows = _generate_windows(0.0, 100.0, 3, 0.7, expanding=False, purge_gap=0.0)
        assert len(windows) > 0

    def test_purge_gap(self):
        windows = _generate_windows(0.0, 100.0, 5, 0.7, expanding=True, purge_gap=5.0)
        for w in windows:
            assert w.test_start >= w.train_end + 5.0

    def test_empty_duration(self):
        windows = _generate_windows(10.0, 10.0, 5, 0.7, expanding=True, purge_gap=0.0)
        assert len(windows) == 0


class TestFilterData:
    def test_filter_dicts(self):
        data = [
            {"timestamp": 1.0, "price": 0.5},
            {"timestamp": 5.0, "price": 0.6},
            {"timestamp": 10.0, "price": 0.7},
        ]
        filtered = _filter_data_by_time(data, 2.0, 8.0)
        assert len(filtered) == 1
        assert filtered[0]["timestamp"] == 5.0

    def test_filter_ticks(self):
        ticks = [Tick(1.0, 0.5), Tick(5.0, 0.6), Tick(10.0, 0.7)]
        filtered = _filter_data_by_time(ticks, 2.0, 8.0)
        assert len(filtered) == 1
        assert filtered[0].timestamp == 5.0

    def test_filter_multi_feed(self):
        data = {
            "feed1": [
                {"timestamp": 1.0, "price": 0.5},
                {"timestamp": 5.0, "price": 0.6},
            ]
        }
        filtered = _filter_data_by_time(data, 2.0, 8.0)
        assert "feed1" in filtered
        assert len(filtered["feed1"]) == 1


class TestWalkForward:
    def test_basic_walk_forward(self):
        """End-to-end test with simple constant pipeline."""
        data = [
            {"timestamp": float(i), "price": 0.50}
            for i in range(1, 101)
        ]

        def pipeline_factory(params):
            spread = params.get("spread", 0.04)
            return [lambda ctx, s=spread: Quote(0.50 - s / 2, 0.50 + s / 2, 5)]

        result = walk_forward(
            data=data,
            pipeline_factory=pipeline_factory,
            param_grid={"spread": [0.04, 0.06]},
            n_splits=3,
            train_ratio=0.7,
        )

        assert isinstance(result, WalkForwardResult)
        assert len(result.windows) == 3
        assert len(result.best_params_per_window) == 3
        assert len(result.test_results) == 3

    def test_single_param_combo(self):
        """With only one parameter combo, that combo is always best."""
        data = [
            {"timestamp": float(i), "price": 0.50}
            for i in range(1, 51)
        ]

        def pipeline_factory(params):
            return [lambda ctx: Quote(0.48, 0.52, 5)]

        result = walk_forward(
            data=data,
            pipeline_factory=pipeline_factory,
            param_grid={"size": [5]},
            n_splits=2,
        )

        assert len(result.best_params_per_window) == 2
        for params in result.best_params_per_window:
            assert params["size"] == 5

    def test_empty_data(self):
        def pipeline_factory(params):
            return [lambda ctx: Quote(0.48, 0.52, 5)]

        result = walk_forward(
            data=[],
            pipeline_factory=pipeline_factory,
            param_grid={"size": [5, 10]},
            n_splits=3,
        )
        assert len(result.windows) == 0

    def test_aggregate_equity_continuity(self):
        """Aggregate equity should chain from previous window's end."""
        data = [
            {"timestamp": float(i), "price": 0.50}
            for i in range(1, 51)
        ]

        def pipeline_factory(params):
            return [lambda ctx: Quote(0.48, 0.52, 5)]

        result = walk_forward(
            data=data,
            pipeline_factory=pipeline_factory,
            param_grid={"size": [5]},
            n_splits=2,
            initial_capital=1000.0,
        )

        if result.aggregate_equity:
            # First equity should start near initial capital
            assert abs(result.aggregate_equity[0][1] - 1000.0) < 100.0
