"""Tests for optimal stopping via Longstaff-Schwartz."""

import math
import pytest

from horizon._horizon import (
    StopPolicy,
    generate_gbm_paths,
    longstaff_schwartz,
    should_exit,
)


# ===========================================================================
# generate_gbm_paths
# ===========================================================================


class TestGenerateGBMPaths:
    def test_basic(self):
        paths = generate_gbm_paths(100.0, 0.05, 0.2, 1.0 / 252.0, 50, 100, 42)
        assert len(paths) == 100
        assert len(paths[0]) == 50

    def test_all_start_at_s0(self):
        paths = generate_gbm_paths(100.0, 0.0, 0.2, 0.01, 20, 50, 42)
        for p in paths:
            assert p[0] == pytest.approx(100.0)

    def test_all_positive(self):
        paths = generate_gbm_paths(100.0, -0.1, 0.3, 0.01, 50, 100, 42)
        for p in paths:
            for v in p:
                assert v > 0.0

    def test_deterministic(self):
        p1 = generate_gbm_paths(100.0, 0.05, 0.2, 0.01, 10, 10, 42)
        p2 = generate_gbm_paths(100.0, 0.05, 0.2, 0.01, 10, 10, 42)
        for a, b in zip(p1, p2):
            for x, y in zip(a, b):
                assert x == pytest.approx(y)

    def test_different_seeds(self):
        p1 = generate_gbm_paths(100.0, 0.05, 0.2, 0.01, 10, 5, 42)
        p2 = generate_gbm_paths(100.0, 0.05, 0.2, 0.01, 10, 5, 99)
        differs = False
        for a, b in zip(p1, p2):
            for x, y in zip(a, b):
                if abs(x - y) > 1e-10:
                    differs = True
                    break
        assert differs

    def test_zero_s0_raises(self):
        with pytest.raises(ValueError):
            generate_gbm_paths(0.0, 0.05, 0.2, 0.01, 10, 10, 42)

    def test_negative_s0_raises(self):
        with pytest.raises(ValueError):
            generate_gbm_paths(-100.0, 0.05, 0.2, 0.01, 10, 10, 42)

    def test_nan_s0_raises(self):
        with pytest.raises(ValueError):
            generate_gbm_paths(float("nan"), 0.05, 0.2, 0.01, 10, 10, 42)

    def test_zero_sigma_raises(self):
        with pytest.raises(ValueError):
            generate_gbm_paths(100.0, 0.05, 0.0, 0.01, 10, 10, 42)

    def test_negative_sigma_raises(self):
        with pytest.raises(ValueError):
            generate_gbm_paths(100.0, 0.05, -0.2, 0.01, 10, 10, 42)

    def test_zero_dt_raises(self):
        with pytest.raises(ValueError):
            generate_gbm_paths(100.0, 0.05, 0.2, 0.0, 10, 10, 42)

    def test_too_few_steps_raises(self):
        with pytest.raises(ValueError):
            generate_gbm_paths(100.0, 0.05, 0.2, 0.01, 1, 10, 42)

    def test_zero_paths_ok(self):
        paths = generate_gbm_paths(100.0, 0.05, 0.2, 0.01, 10, 0, 42)
        assert len(paths) == 0

    def test_nan_mu_raises(self):
        with pytest.raises(ValueError):
            generate_gbm_paths(100.0, float("nan"), 0.2, 0.01, 10, 10, 42)

    def test_large_volatility(self):
        paths = generate_gbm_paths(100.0, 0.0, 1.0, 0.01, 20, 10, 42)
        assert len(paths) == 10
        for p in paths:
            for v in p:
                assert v > 0.0

    def test_negative_drift(self):
        paths = generate_gbm_paths(100.0, -0.5, 0.2, 0.01, 50, 10, 42)
        # With strong negative drift, average final price should be below S0
        finals = [p[-1] for p in paths]
        avg = sum(finals) / len(finals)
        assert avg < 100.0

    def test_path_dimensions(self):
        paths = generate_gbm_paths(50.0, 0.1, 0.3, 0.01, 30, 20, 42)
        assert len(paths) == 20
        for p in paths:
            assert len(p) == 30


# ===========================================================================
# longstaff_schwartz
# ===========================================================================


class TestLongstaffSchwartz:
    def test_basic(self):
        paths = generate_gbm_paths(100.0, -0.1, 0.3, 1.0 / 252.0, 20, 500, 42)
        policy = longstaff_schwartz(paths, 100.0)
        assert isinstance(policy, StopPolicy)

    def test_result_fields(self):
        paths = generate_gbm_paths(100.0, -0.1, 0.3, 0.01, 20, 100, 42)
        policy = longstaff_schwartz(paths, 100.0)
        assert hasattr(policy, "exercise_boundary")
        assert hasattr(policy, "expected_payoff")
        assert hasattr(policy, "optimal_stop_time")

    def test_boundary_length(self):
        paths = generate_gbm_paths(100.0, -0.05, 0.2, 0.01, 20, 100, 42)
        policy = longstaff_schwartz(paths, 100.0)
        assert len(policy.exercise_boundary) == 20

    def test_payoff_non_negative(self):
        paths = generate_gbm_paths(100.0, -0.1, 0.3, 0.01, 20, 500, 42)
        policy = longstaff_schwartz(paths, 100.0)
        assert policy.expected_payoff >= 0.0

    def test_stop_time_in_range(self):
        paths = generate_gbm_paths(100.0, -0.1, 0.3, 0.01, 20, 500, 42)
        policy = longstaff_schwartz(paths, 100.0)
        assert 0.0 <= policy.optimal_stop_time <= 1.0

    def test_custom_n_basis(self):
        paths = generate_gbm_paths(100.0, -0.1, 0.3, 0.01, 20, 100, 42)
        policy = longstaff_schwartz(paths, 100.0, n_basis=3)
        assert len(policy.exercise_boundary) == 20

    def test_empty_paths_raises(self):
        with pytest.raises(ValueError):
            longstaff_schwartz([], 100.0)

    def test_too_few_paths_raises(self):
        paths = generate_gbm_paths(100.0, 0.0, 0.2, 0.01, 10, 5, 42)
        with pytest.raises(ValueError):
            longstaff_schwartz(paths, 100.0)

    def test_too_few_steps_raises(self):
        paths = [[100.0, 95.0] for _ in range(50)]
        with pytest.raises(ValueError):
            longstaff_schwartz(paths, 100.0)

    def test_zero_strike_raises(self):
        paths = generate_gbm_paths(100.0, 0.0, 0.2, 0.01, 10, 50, 42)
        with pytest.raises(ValueError):
            longstaff_schwartz(paths, 0.0)

    def test_negative_strike_raises(self):
        paths = generate_gbm_paths(100.0, 0.0, 0.2, 0.01, 10, 50, 42)
        with pytest.raises(ValueError):
            longstaff_schwartz(paths, -100.0)

    def test_n_basis_too_small_raises(self):
        paths = generate_gbm_paths(100.0, 0.0, 0.2, 0.01, 10, 50, 42)
        with pytest.raises(ValueError):
            longstaff_schwartz(paths, 100.0, n_basis=1)

    def test_n_basis_too_large_raises(self):
        paths = generate_gbm_paths(100.0, 0.0, 0.2, 0.01, 10, 50, 42)
        with pytest.raises(ValueError):
            longstaff_schwartz(paths, 100.0, n_basis=21)

    def test_repr(self):
        paths = generate_gbm_paths(100.0, -0.1, 0.3, 0.01, 20, 100, 42)
        policy = longstaff_schwartz(paths, 100.0)
        r = repr(policy)
        assert "StopPolicy" in r

    def test_boundary_values_finite(self):
        paths = generate_gbm_paths(100.0, -0.1, 0.3, 0.01, 20, 200, 42)
        policy = longstaff_schwartz(paths, 100.0)
        for b in policy.exercise_boundary:
            assert math.isfinite(b)

    def test_more_paths_stable_payoff(self):
        paths1 = generate_gbm_paths(100.0, -0.1, 0.3, 0.01, 15, 100, 42)
        paths2 = generate_gbm_paths(100.0, -0.1, 0.3, 0.01, 15, 1000, 42)
        pol1 = longstaff_schwartz(paths1, 100.0)
        pol2 = longstaff_schwartz(paths2, 100.0)
        # Both should produce finite payoffs
        assert math.isfinite(pol1.expected_payoff)
        assert math.isfinite(pol2.expected_payoff)


# ===========================================================================
# should_exit
# ===========================================================================


class TestShouldExit:
    def test_below_boundary_exits(self):
        boundary = [90.0, 85.0, 80.0, 75.0]
        assert should_exit(70.0, boundary, 2)

    def test_above_boundary_no_exit(self):
        boundary = [90.0, 85.0, 80.0, 75.0]
        assert not should_exit(95.0, boundary, 0)

    def test_at_boundary_exits(self):
        boundary = [90.0, 85.0, 80.0, 75.0]
        assert should_exit(85.0, boundary, 1)

    def test_past_horizon_exits(self):
        boundary = [90.0, 85.0]
        assert should_exit(100.0, boundary, 5)

    def test_empty_boundary_exits(self):
        assert should_exit(100.0, [], 0)

    def test_first_step(self):
        boundary = [95.0, 90.0, 85.0]
        assert should_exit(94.0, boundary, 0)
        assert not should_exit(96.0, boundary, 0)

    def test_last_step(self):
        boundary = [95.0, 90.0, 85.0]
        assert should_exit(84.0, boundary, 2)
        assert not should_exit(86.0, boundary, 2)

    def test_returns_bool(self):
        boundary = [90.0]
        result = should_exit(85.0, boundary, 0)
        assert isinstance(result, bool)

    def test_all_steps_below(self):
        boundary = [100.0, 100.0, 100.0]
        for step in range(3):
            assert should_exit(50.0, boundary, step)

    def test_all_steps_above(self):
        boundary = [50.0, 50.0, 50.0]
        for step in range(3):
            assert not should_exit(100.0, boundary, step)
