"""Tests for Bayesian optimization module."""

from __future__ import annotations

import math
import random

import pytest
from horizon._horizon import Quote
from horizon.bayesian_opt import (
    BayesianOptResult,
    bayesian_optimize,
    _rbf_kernel,
    _kernel_matrix,
    _cholesky,
    _cholesky_solve,
    _GP,
    _expected_improvement,
    _normal_cdf,
    _normal_pdf,
    _normalize_params,
    _denormalize_vec,
    _estimate_length_scale,
    _maximize_acquisition,
)


# ---------------------------------------------------------------------------
# GP kernel tests
# ---------------------------------------------------------------------------


class TestRBFKernel:
    def test_identical_points(self):
        """k(x, x) should be 1.0 for RBF kernel."""
        x = [0.5, 0.3, 0.7]
        assert abs(_rbf_kernel(x, x, length_scale=1.0) - 1.0) < 1e-12

    def test_distant_points(self):
        """Distant points should have kernel value close to 0."""
        x1 = [0.0, 0.0]
        x2 = [100.0, 100.0]
        k = _rbf_kernel(x1, x2, length_scale=1.0)
        assert k < 1e-10

    def test_symmetry(self):
        """k(x, y) == k(y, x)."""
        x = [0.2, 0.8]
        y = [0.7, 0.3]
        assert abs(_rbf_kernel(x, y, 0.5) - _rbf_kernel(y, x, 0.5)) < 1e-15

    def test_length_scale_effect(self):
        """Larger length scale => higher kernel value for the same pair."""
        x = [0.0]
        y = [1.0]
        k_small = _rbf_kernel(x, y, length_scale=0.1)
        k_large = _rbf_kernel(x, y, length_scale=10.0)
        assert k_large > k_small

    def test_1d_known_value(self):
        """Check a known analytical value."""
        # k([0], [1], l=1) = exp(-0.5)
        expected = math.exp(-0.5)
        got = _rbf_kernel([0.0], [1.0], length_scale=1.0)
        assert abs(got - expected) < 1e-12


class TestKernelMatrix:
    def test_shape(self):
        X = [[0.1], [0.5], [0.9]]
        K = _kernel_matrix(X, length_scale=1.0)
        assert len(K) == 3
        assert all(len(row) == 3 for row in K)

    def test_symmetry(self):
        X = [[0.1, 0.2], [0.5, 0.6], [0.9, 0.1]]
        K = _kernel_matrix(X, length_scale=0.5)
        for i in range(3):
            for j in range(3):
                assert abs(K[i][j] - K[j][i]) < 1e-15

    def test_positive_diagonal(self):
        X = [[0.1], [0.5], [0.9]]
        K = _kernel_matrix(X, length_scale=1.0, noise=1e-6)
        for i in range(3):
            assert K[i][i] > 1.0  # 1.0 from kernel + noise


# ---------------------------------------------------------------------------
# Cholesky / solve tests
# ---------------------------------------------------------------------------


class TestCholesky:
    def test_identity(self):
        A = [[1.0, 0.0], [0.0, 1.0]]
        L = _cholesky(A)
        assert abs(L[0][0] - 1.0) < 1e-12
        assert abs(L[1][1] - 1.0) < 1e-12
        assert abs(L[0][1]) < 1e-12
        assert abs(L[1][0]) < 1e-12

    def test_known_matrix(self):
        # A = [[4, 2], [2, 3]]  =>  L = [[2, 0], [1, sqrt(2)]]
        A = [[4.0, 2.0], [2.0, 3.0]]
        L = _cholesky(A)
        assert abs(L[0][0] - 2.0) < 1e-10
        assert abs(L[1][0] - 1.0) < 1e-10
        assert abs(L[1][1] - math.sqrt(2.0)) < 1e-10

    def test_jitter_recovery(self):
        """Near-singular matrix should still produce a valid L with jitter."""
        A = [[1.0, 1.0 - 1e-12], [1.0 - 1e-12, 1.0]]
        L = _cholesky(A)
        # Just check it didn't crash and diagonal is positive
        for i in range(2):
            assert L[i][i] > 0


class TestCholeskySolve:
    def test_simple_system(self):
        # Solve [[4, 2], [2, 3]] x = [8, 7]
        # x = [1, 1] (by inspection: 4+2=6 nope... 4*1+2*1.333=6.666 nope)
        # Actually: 4x + 2y = 8, 2x + 3y = 7 => x=1, y=2? 4+4=8 yes, 2+6=8 nope
        # 4x+2y=8, 2x+3y=7 => from first: x=2-y/2, sub: 4-y+3y=7 => 2y=3 => y=1.5, x=1.25
        A = [[4.0, 2.0], [2.0, 3.0]]
        b = [8.0, 7.0]
        L = _cholesky(A)
        x = _cholesky_solve(L, b)
        assert abs(x[0] - 1.25) < 1e-8
        assert abs(x[1] - 1.5) < 1e-8


# ---------------------------------------------------------------------------
# GP prediction tests
# ---------------------------------------------------------------------------


class TestGP:
    def test_predict_at_observed_point(self):
        """GP should interpolate through training data (low noise)."""
        gp = _GP(length_scale=0.3, noise=1e-8)
        X = [[0.0], [0.5], [1.0]]
        y = [1.0, 2.0, 3.0]
        gp.fit(X, y)

        mu, std = gp.predict([0.5])
        assert abs(mu - 2.0) < 0.05  # should be very close to observed value
        assert std < 0.1  # low uncertainty at observed point

    def test_uncertainty_far_from_data(self):
        """Uncertainty should be higher far from observed data."""
        gp = _GP(length_scale=0.3, noise=1e-6)
        X = [[0.0], [0.1]]
        y = [1.0, 1.1]
        gp.fit(X, y)

        _, std_near = gp.predict([0.05])
        _, std_far = gp.predict([5.0])
        assert std_far > std_near

    def test_predict_quadratic(self):
        """GP should roughly learn a simple quadratic in 1D."""
        gp = _GP(length_scale=0.5, noise=1e-6)
        X = [[x / 10.0] for x in range(11)]  # 0.0 to 1.0
        y = [(x / 10.0) ** 2 for x in range(11)]  # y = x^2
        gp.fit(X, y)

        mu, _ = gp.predict([0.5])
        assert abs(mu - 0.25) < 0.1  # should be close to 0.5^2 = 0.25

    def test_predict_empty(self):
        """GP with no data should return prior mean and unit std."""
        gp = _GP()
        mu, std = gp.predict([0.5])
        assert math.isfinite(mu)
        assert std == 1.0

    def test_all_outputs_finite(self):
        """All GP predictions must be finite."""
        gp = _GP(length_scale=0.3, noise=1e-6)
        rng = random.Random(42)
        X = [[rng.random()] for _ in range(20)]
        y = [rng.gauss(0, 1) for _ in range(20)]
        gp.fit(X, y)

        for _ in range(50):
            x = [rng.random()]
            mu, std = gp.predict(x)
            assert math.isfinite(mu), f"mu={mu} is not finite"
            assert math.isfinite(std), f"std={std} is not finite"
            assert std >= 0.0


# ---------------------------------------------------------------------------
# Acquisition function tests
# ---------------------------------------------------------------------------


class TestExpectedImprovement:
    def test_zero_std(self):
        """EI should be 0 when std is 0 (no uncertainty)."""
        assert _expected_improvement(mu=1.0, std=0.0, f_best=0.5) == 0.0

    def test_positive_when_mu_above_best(self):
        """EI should be positive when predicted mean exceeds current best."""
        ei = _expected_improvement(mu=2.0, std=0.5, f_best=1.0)
        assert ei > 0.0

    def test_positive_even_below_best_with_uncertainty(self):
        """High uncertainty can produce positive EI even when mu < f_best."""
        ei = _expected_improvement(mu=0.8, std=2.0, f_best=1.0)
        assert ei > 0.0

    def test_higher_std_higher_ei(self):
        """More uncertainty should generally increase EI (exploration)."""
        ei_low = _expected_improvement(mu=1.0, std=0.1, f_best=1.0)
        ei_high = _expected_improvement(mu=1.0, std=1.0, f_best=1.0)
        assert ei_high > ei_low

    def test_non_negative(self):
        """EI should never be negative."""
        rng = random.Random(123)
        for _ in range(100):
            mu = rng.gauss(0, 5)
            std = abs(rng.gauss(0, 2))
            f_best = rng.gauss(0, 5)
            ei = _expected_improvement(mu, std, f_best)
            assert ei >= 0.0, f"EI={ei} < 0 for mu={mu}, std={std}, f_best={f_best}"


class TestNormalDistribution:
    def test_cdf_at_zero(self):
        assert abs(_normal_cdf(0.0) - 0.5) < 1e-10

    def test_cdf_large_positive(self):
        assert _normal_cdf(10.0) > 0.999

    def test_cdf_large_negative(self):
        assert _normal_cdf(-10.0) < 0.001

    def test_pdf_at_zero(self):
        expected = 1.0 / math.sqrt(2.0 * math.pi)
        assert abs(_normal_pdf(0.0) - expected) < 1e-12

    def test_pdf_symmetry(self):
        assert abs(_normal_pdf(1.5) - _normal_pdf(-1.5)) < 1e-15


# ---------------------------------------------------------------------------
# Normalization tests
# ---------------------------------------------------------------------------


class TestNormalization:
    def test_normalize_round_trip(self):
        space = {"a": (0.0, 10.0), "b": (-5.0, 5.0)}
        params = {"a": 3.0, "b": 2.0}
        vec = _normalize_params(params, space)
        recovered = _denormalize_vec(vec, space)
        assert abs(recovered["a"] - 3.0) < 1e-12
        assert abs(recovered["b"] - 2.0) < 1e-12

    def test_normalize_bounds(self):
        space = {"x": (0.0, 1.0)}
        assert abs(_normalize_params({"x": 0.0}, space)[0]) < 1e-12
        assert abs(_normalize_params({"x": 1.0}, space)[0] - 1.0) < 1e-12

    def test_zero_range(self):
        space = {"x": (5.0, 5.0)}
        vec = _normalize_params({"x": 5.0}, space)
        assert vec == [0.5]


class TestEstimateLengthScale:
    def test_basic(self):
        space = {"a": (0.0, 3.0), "b": (0.0, 6.0)}
        ls = _estimate_length_scale(space)
        assert ls > 0
        # average range = 4.5, ls ~ 4.5/3 = 1.5
        assert abs(ls - 1.5) < 0.01

    def test_empty(self):
        assert _estimate_length_scale({}) == 1.0


# ---------------------------------------------------------------------------
# Acquisition maximizer test
# ---------------------------------------------------------------------------


class TestMaximizeAcquisition:
    def test_returns_valid_params(self):
        gp = _GP(length_scale=0.3, noise=1e-4)
        gp.fit([[0.2], [0.8]], [1.0, 0.5])
        space = {"x": (0.0, 1.0)}
        params = _maximize_acquisition(gp, f_best=1.0, search_space=space, rng=random.Random(42))
        assert "x" in params
        assert 0.0 <= params["x"] <= 1.0


# ---------------------------------------------------------------------------
# Full optimization tests (end-to-end with backtest)
# ---------------------------------------------------------------------------


class TestBayesianOptResult:
    def test_default(self):
        r = BayesianOptResult()
        assert r.best_params == {}
        assert r.best_score == float("-inf")
        assert r.all_trials == []
        assert r.n_trials == 0
        assert r.convergence == []


class TestBayesianOptimize:
    """End-to-end tests using a simple constant pipeline."""

    @staticmethod
    def _make_data(n: int = 50) -> list[dict]:
        """Create simple price data for backtesting."""
        return [
            {"timestamp": float(i), "price": 0.50}
            for i in range(1, n + 1)
        ]

    @staticmethod
    def _simple_factory(params: dict) -> list:
        """Pipeline that quotes around 0.50 with parameterized spread."""
        spread = params.get("spread", 0.04)
        half = spread / 2.0
        bid = max(0.01, 0.50 - half)
        ask = min(0.99, 0.50 + half)

        def pipeline_fn(ctx):
            if bid >= ask:
                return None
            return Quote(bid, ask, 5)

        return [pipeline_fn]

    def test_basic_optimization(self):
        data = self._make_data()
        result = bayesian_optimize(
            data=data,
            pipeline_factory=self._simple_factory,
            search_space={"spread": (0.02, 0.20)},
            objective="total_return",
            n_trials=12,
            n_initial=6,
            seed=42,
        )
        assert isinstance(result, BayesianOptResult)
        assert result.n_trials == 12
        assert len(result.all_trials) == 12
        assert len(result.convergence) == 12
        assert "spread" in result.best_params
        assert math.isfinite(result.best_score)

    def test_convergence_monotonic(self):
        """Convergence trace should be monotonically non-decreasing."""
        data = self._make_data()
        result = bayesian_optimize(
            data=data,
            pipeline_factory=self._simple_factory,
            search_space={"spread": (0.02, 0.20)},
            objective="total_return",
            n_trials=15,
            n_initial=8,
            seed=99,
        )
        for i in range(1, len(result.convergence)):
            assert result.convergence[i] >= result.convergence[i - 1]

    def test_single_trial(self):
        """n_trials=1 should work and return exactly one trial."""
        data = self._make_data()
        result = bayesian_optimize(
            data=data,
            pipeline_factory=self._simple_factory,
            search_space={"spread": (0.02, 0.10)},
            n_trials=1,
            n_initial=1,
            seed=7,
        )
        assert result.n_trials == 1
        assert len(result.all_trials) == 1

    def test_zero_trials(self):
        data = self._make_data()
        result = bayesian_optimize(
            data=data,
            pipeline_factory=self._simple_factory,
            search_space={"spread": (0.02, 0.10)},
            n_trials=0,
        )
        assert result.n_trials == 0

    def test_seed_reproducibility(self):
        """Same seed should produce identical results."""
        data = self._make_data()
        kwargs = dict(
            data=data,
            pipeline_factory=self._simple_factory,
            search_space={"spread": (0.02, 0.20)},
            n_trials=10,
            n_initial=5,
            seed=12345,
            objective="total_return",
        )
        r1 = bayesian_optimize(**kwargs)
        r2 = bayesian_optimize(**kwargs)
        assert r1.best_params == r2.best_params
        assert r1.best_score == r2.best_score
        assert r1.convergence == r2.convergence

    def test_different_seeds_differ(self):
        """Different seeds should (almost certainly) produce different trials."""
        data = self._make_data()
        base = dict(
            data=data,
            pipeline_factory=self._simple_factory,
            search_space={"spread": (0.02, 0.20)},
            n_trials=8,
            n_initial=8,
            objective="total_return",
        )
        r1 = bayesian_optimize(**base, seed=1)
        r2 = bayesian_optimize(**base, seed=999)
        # At least one trial should differ
        params_1 = [t["params"]["spread"] for t in r1.all_trials]
        params_2 = [t["params"]["spread"] for t in r2.all_trials]
        assert params_1 != params_2

    def test_single_param(self):
        """Optimization with a single parameter dimension."""
        data = self._make_data()
        result = bayesian_optimize(
            data=data,
            pipeline_factory=self._simple_factory,
            search_space={"spread": (0.02, 0.10)},
            n_trials=10,
            n_initial=5,
            seed=42,
        )
        assert result.n_trials == 10
        assert "spread" in result.best_params

    def test_multi_param(self):
        """Optimization with multiple parameters."""
        data = self._make_data()

        def factory(params):
            spread = params.get("spread", 0.04)
            size = params.get("size", 5.0)
            half = spread / 2.0
            bid = max(0.01, 0.50 - half)
            ask = min(0.99, 0.50 + half)

            def fn(ctx):
                if bid >= ask:
                    return None
                return Quote(bid, ask, max(1.0, size))

            return [fn]

        result = bayesian_optimize(
            data=data,
            pipeline_factory=factory,
            search_space={"spread": (0.02, 0.20), "size": (1.0, 20.0)},
            n_trials=12,
            n_initial=6,
            seed=42,
        )
        assert result.n_trials == 12
        assert "spread" in result.best_params
        assert "size" in result.best_params

    def test_best_score_matches_best_trial(self):
        """best_score should equal the max score in all_trials."""
        data = self._make_data()
        result = bayesian_optimize(
            data=data,
            pipeline_factory=self._simple_factory,
            search_space={"spread": (0.02, 0.20)},
            n_trials=10,
            n_initial=5,
            seed=42,
        )
        max_trial_score = max(t["score"] for t in result.all_trials)
        assert abs(result.best_score - max_trial_score) < 1e-12

    def test_params_within_bounds(self):
        """All sampled parameters should lie within the search space."""
        data = self._make_data()
        result = bayesian_optimize(
            data=data,
            pipeline_factory=self._simple_factory,
            search_space={"spread": (0.02, 0.20)},
            n_trials=15,
            n_initial=8,
            seed=42,
        )
        for trial in result.all_trials:
            s = trial["params"]["spread"]
            assert 0.02 <= s <= 0.20 + 1e-9, f"spread={s} out of bounds"

    def test_all_scores_finite(self):
        """Every trial score should be a finite number."""
        data = self._make_data()
        result = bayesian_optimize(
            data=data,
            pipeline_factory=self._simple_factory,
            search_space={"spread": (0.02, 0.20)},
            n_trials=10,
            n_initial=5,
            seed=42,
        )
        for trial in result.all_trials:
            assert math.isfinite(trial["score"]), f"Non-finite score: {trial['score']}"

    def test_n_initial_capped_to_n_trials(self):
        """When n_initial > n_trials, it should be capped."""
        data = self._make_data()
        result = bayesian_optimize(
            data=data,
            pipeline_factory=self._simple_factory,
            search_space={"spread": (0.02, 0.10)},
            n_trials=3,
            n_initial=100,  # way more than n_trials
            seed=42,
        )
        assert result.n_trials == 3

    def test_factory_exception_handled(self):
        """If pipeline_factory or backtest raises, score should be 0."""
        data = self._make_data()

        call_count = 0

        def bad_factory(params):
            nonlocal call_count
            call_count += 1
            if call_count % 2 == 0:
                raise ValueError("intentional failure")
            spread = params.get("spread", 0.04)
            half = spread / 2.0
            return [lambda ctx: Quote(0.50 - half, 0.50 + half, 5)]

        result = bayesian_optimize(
            data=data,
            pipeline_factory=bad_factory,
            search_space={"spread": (0.02, 0.10)},
            n_trials=6,
            n_initial=6,
            seed=42,
        )
        assert result.n_trials == 6
        # Some scores should be 0 from the exceptions
        scores = [t["score"] for t in result.all_trials]
        assert 0.0 in scores

    def test_gp_guided_trials_evaluated(self):
        """Verify GP-guided phase actually runs (n_trials > n_initial)."""
        data = self._make_data(80)
        result = bayesian_optimize(
            data=data,
            pipeline_factory=self._simple_factory,
            search_space={"spread": (0.02, 0.20)},
            n_trials=15,
            n_initial=5,
            seed=42,
        )
        # All 15 trials should have been evaluated
        assert result.n_trials == 15
        assert len(result.all_trials) == 15
