"""Tests for entropy pooling (Meucci)."""

import math
import pytest

from horizon._horizon import (
    EntropyPoolResult,
    ViewConstraint,
    entropy_pool,
    posterior_covariance,
    posterior_mean,
)


# ===========================================================================
# Helper data
# ===========================================================================

def _uniform_probs(n=5):
    """Uniform prior probabilities."""
    return [1.0 / n] * n


def _simple_scenarios(n=5, dim=2):
    """Simple scenarios: each scenario is a vector of asset returns."""
    scenarios = []
    for i in range(n):
        scenario = [0.01 * (i + 1) * (j + 1) for j in range(dim)]
        scenarios.append(scenario)
    return scenarios


def _diverse_scenarios(n=10, dim=3):
    """More diverse scenarios for better testing."""
    import random
    rng = random.Random(42)
    scenarios = []
    for _ in range(n):
        scenario = [rng.gauss(0, 0.05) for _ in range(dim)]
        scenarios.append(scenario)
    return scenarios


# ===========================================================================
# ViewConstraint
# ===========================================================================


class TestViewConstraint:
    def test_creation(self):
        vc = ViewConstraint([1.0, 0.0], 0.05)
        assert vc is not None

    def test_fields(self):
        vc = ViewConstraint([1.0, -1.0], 0.02)
        assert hasattr(vc, "coefficients")
        assert hasattr(vc, "bound")
        assert hasattr(vc, "is_equality")

    def test_coefficients_accessible(self):
        vc = ViewConstraint([0.5, 0.5], 0.03)
        assert vc.coefficients == [0.5, 0.5]

    def test_bound_accessible(self):
        vc = ViewConstraint([1.0, 0.0], 0.05)
        assert vc.bound == pytest.approx(0.05)

    def test_negative_bound(self):
        vc = ViewConstraint([1.0, 0.0], -0.03)
        assert vc.bound == pytest.approx(-0.03)

    def test_zero_bound(self):
        vc = ViewConstraint([1.0, -1.0], 0.0)
        assert vc.bound == pytest.approx(0.0)

    def test_is_equality_default(self):
        vc = ViewConstraint([1.0, 0.0], 0.05)
        assert isinstance(vc.is_equality, bool)

    def test_is_equality_true(self):
        vc = ViewConstraint([1.0, 0.0], 0.05, True)
        assert vc.is_equality is True

    def test_repr(self):
        vc = ViewConstraint([1.0, 0.0], 0.05)
        r = repr(vc)
        assert "ViewConstraint" in r


# ===========================================================================
# entropy_pool
# ===========================================================================


class TestEntropyPool:
    def test_basic(self):
        probs = _uniform_probs(5)
        scenarios = _simple_scenarios(5, 2)
        view = ViewConstraint([1.0, 0.0], 0.03)
        result = entropy_pool(probs, scenarios, [view])
        assert isinstance(result, EntropyPoolResult)

    def test_result_fields(self):
        probs = _uniform_probs(5)
        scenarios = _simple_scenarios(5, 2)
        view = ViewConstraint([1.0, 0.0], 0.03)
        result = entropy_pool(probs, scenarios, [view])
        assert hasattr(result, "posterior_probs")
        assert hasattr(result, "prior_probs")
        assert hasattr(result, "relative_entropy")
        assert hasattr(result, "effective_scenarios")

    def test_posterior_probs_sum_to_one(self):
        probs = _uniform_probs(10)
        scenarios = _simple_scenarios(10, 2)
        view = ViewConstraint([1.0, 0.0], 0.05)
        result = entropy_pool(probs, scenarios, [view])
        total = sum(result.posterior_probs)
        assert total == pytest.approx(1.0, abs=1e-6)

    def test_posterior_probs_non_negative(self):
        probs = _uniform_probs(5)
        scenarios = _simple_scenarios(5, 2)
        view = ViewConstraint([1.0, 0.0], 0.03)
        result = entropy_pool(probs, scenarios, [view])
        for p in result.posterior_probs:
            assert p >= -1e-10

    def test_no_views_raises(self):
        probs = _uniform_probs(5)
        scenarios = _simple_scenarios(5, 2)
        with pytest.raises(ValueError):
            entropy_pool(probs, scenarios, [])

    def test_multiple_views(self):
        probs = _uniform_probs(10)
        scenarios = _simple_scenarios(10, 3)
        v1 = ViewConstraint([1.0, 0.0, 0.0], 0.05)
        v2 = ViewConstraint([0.0, 1.0, -1.0], 0.01)
        result = entropy_pool(probs, scenarios, [v1, v2])
        total = sum(result.posterior_probs)
        assert total == pytest.approx(1.0, abs=1e-6)

    def test_relative_entropy_non_negative(self):
        probs = _uniform_probs(5)
        scenarios = _simple_scenarios(5, 2)
        view = ViewConstraint([1.0, 0.0], 0.03)
        result = entropy_pool(probs, scenarios, [view])
        assert result.relative_entropy >= -1e-10

    def test_effective_scenarios_bounded(self):
        probs = _uniform_probs(10)
        scenarios = _simple_scenarios(10, 2)
        view = ViewConstraint([1.0, 0.0], 0.05)
        result = entropy_pool(probs, scenarios, [view])
        assert 0.0 < result.effective_scenarios <= 10.0

    def test_prior_probs_preserved(self):
        probs = _uniform_probs(5)
        scenarios = _simple_scenarios(5, 2)
        view = ViewConstraint([1.0, 0.0], 0.03)
        result = entropy_pool(probs, scenarios, [view])
        for p_prior, p_stored in zip(probs, result.prior_probs):
            assert p_prior == pytest.approx(p_stored, abs=1e-10)

    def test_posterior_probs_length(self):
        n = 8
        probs = _uniform_probs(n)
        scenarios = _simple_scenarios(n, 2)
        view = ViewConstraint([1.0, 0.0], 0.04)
        result = entropy_pool(probs, scenarios, [view])
        assert len(result.posterior_probs) == n

    def test_repr(self):
        probs = _uniform_probs(5)
        scenarios = _simple_scenarios(5, 2)
        view = ViewConstraint([1.0, 0.0], 0.03)
        result = entropy_pool(probs, scenarios, [view])
        r = repr(result)
        assert "EntropyPoolResult" in r

    def test_diverse_scenarios(self):
        probs = _uniform_probs(10)
        scenarios = _diverse_scenarios(10, 3)
        view = ViewConstraint([1.0, 0.0, 0.0], 0.01)
        result = entropy_pool(probs, scenarios, [view])
        total = sum(result.posterior_probs)
        assert total == pytest.approx(1.0, abs=1e-6)


# ===========================================================================
# posterior_mean
# ===========================================================================


class TestPosteriorMean:
    def test_basic(self):
        probs = _uniform_probs(5)
        scenarios = _simple_scenarios(5, 2)
        mean = posterior_mean(probs, scenarios)
        assert len(mean) == 2
        for m in mean:
            assert math.isfinite(m)

    def test_uniform_probs(self):
        probs = [0.5, 0.5]
        scenarios = [[1.0, 2.0], [3.0, 4.0]]
        mean = posterior_mean(probs, scenarios)
        assert mean[0] == pytest.approx(2.0)
        assert mean[1] == pytest.approx(3.0)

    def test_single_scenario(self):
        probs = [1.0]
        scenarios = [[5.0, 10.0]]
        mean = posterior_mean(probs, scenarios)
        assert mean[0] == pytest.approx(5.0)
        assert mean[1] == pytest.approx(10.0)

    def test_concentrated_probability(self):
        probs = [0.01, 0.01, 0.98]
        scenarios = [[1.0], [2.0], [100.0]]
        mean = posterior_mean(probs, scenarios)
        assert mean[0] == pytest.approx(98.03, abs=0.1)

    def test_three_assets(self):
        probs = _uniform_probs(5)
        scenarios = _simple_scenarios(5, 3)
        mean = posterior_mean(probs, scenarios)
        assert len(mean) == 3


# ===========================================================================
# posterior_covariance
# ===========================================================================


class TestPosteriorCovariance:
    def test_basic(self):
        probs = _uniform_probs(5)
        scenarios = _simple_scenarios(5, 2)
        cov = posterior_covariance(probs, scenarios)
        assert len(cov) == 2
        for row in cov:
            assert len(row) == 2

    def test_diagonal_non_negative(self):
        probs = _uniform_probs(5)
        scenarios = _simple_scenarios(5, 2)
        cov = posterior_covariance(probs, scenarios)
        for i in range(len(cov)):
            assert cov[i][i] >= -1e-10

    def test_symmetric(self):
        probs = _uniform_probs(5)
        scenarios = _simple_scenarios(5, 3)
        cov = posterior_covariance(probs, scenarios)
        for i in range(len(cov)):
            for j in range(len(cov)):
                assert cov[i][j] == pytest.approx(cov[j][i], abs=1e-10)

    def test_single_scenario_zero_variance(self):
        probs = [1.0]
        scenarios = [[5.0, 10.0]]
        cov = posterior_covariance(probs, scenarios)
        for i in range(2):
            for j in range(2):
                assert cov[i][j] == pytest.approx(0.0, abs=1e-10)

    def test_entries_finite(self):
        probs = _uniform_probs(10)
        scenarios = _diverse_scenarios(10, 3)
        cov = posterior_covariance(probs, scenarios)
        for row in cov:
            for val in row:
                assert math.isfinite(val)
