"""Tests for robust portfolio optimization."""

import math
import pytest

from horizon._horizon import (
    RobustPortfolioResult,
    robust_efficient_frontier,
    robust_optimize,
    worst_case_return,
)


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

def _simple_inputs():
    """2-asset simple inputs."""
    mu = [0.10, 0.15]
    sigma = [[0.04, 0.01], [0.01, 0.09]]
    unc = [0.02, 0.03]
    return mu, sigma, unc


def _three_asset_inputs():
    """3-asset inputs."""
    mu = [0.08, 0.12, 0.06]
    sigma = [
        [0.04, 0.01, 0.005],
        [0.01, 0.09, 0.02],
        [0.005, 0.02, 0.025],
    ]
    unc = [0.02, 0.03, 0.01]
    return mu, sigma, unc


def _four_asset_inputs():
    """4-asset inputs."""
    mu = [0.05, 0.10, 0.15, 0.08]
    sigma = [
        [0.04, 0.01, 0.005, 0.002],
        [0.01, 0.09, 0.02, 0.01],
        [0.005, 0.02, 0.16, 0.03],
        [0.002, 0.01, 0.03, 0.04],
    ]
    unc = [0.01, 0.02, 0.04, 0.015]
    return mu, sigma, unc


# ===========================================================================
# robust_optimize
# ===========================================================================


class TestRobustOptimize:
    def test_basic(self):
        mu, sigma, unc = _simple_inputs()
        result = robust_optimize(mu, sigma, unc, 0.5)
        assert isinstance(result, RobustPortfolioResult)

    def test_result_fields(self):
        mu, sigma, unc = _simple_inputs()
        result = robust_optimize(mu, sigma, unc, 0.5)
        assert hasattr(result, "weights")
        assert hasattr(result, "worst_case_return")
        assert hasattr(result, "nominal_return")
        assert hasattr(result, "uncertainty_budget")

    def test_weights_sum_to_one(self):
        mu, sigma, unc = _simple_inputs()
        result = robust_optimize(mu, sigma, unc, 0.5)
        assert sum(result.weights) == pytest.approx(1.0, abs=0.01)

    def test_weights_non_negative_by_default(self):
        mu, sigma, unc = _simple_inputs()
        result = robust_optimize(mu, sigma, unc, 0.5)
        for w in result.weights:
            assert w >= -0.01

    def test_worst_case_leq_nominal(self):
        mu, sigma, unc = _simple_inputs()
        result = robust_optimize(mu, sigma, unc, 0.5)
        assert result.worst_case_return <= result.nominal_return + 1e-6

    def test_uncertainty_budget_stored(self):
        mu, sigma, unc = _simple_inputs()
        result = robust_optimize(mu, sigma, unc, 0.7)
        assert result.uncertainty_budget == pytest.approx(0.7)

    def test_zero_budget_wc_equals_nominal(self):
        mu, sigma, unc = _simple_inputs()
        result = robust_optimize(mu, sigma, unc, 0.0)
        assert result.worst_case_return == pytest.approx(
            result.nominal_return, abs=0.01
        )

    def test_three_assets(self):
        mu, sigma, unc = _three_asset_inputs()
        result = robust_optimize(mu, sigma, unc, 0.3)
        assert len(result.weights) == 3
        assert sum(result.weights) == pytest.approx(1.0, abs=0.01)

    def test_four_assets(self):
        mu, sigma, unc = _four_asset_inputs()
        result = robust_optimize(mu, sigma, unc, 0.5)
        assert len(result.weights) == 4
        assert sum(result.weights) == pytest.approx(1.0, abs=0.01)

    def test_custom_weight_bounds(self):
        mu, sigma, unc = _simple_inputs()
        result = robust_optimize(mu, sigma, unc, 0.5, max_weight=0.6, min_weight=0.1)
        for w in result.weights:
            assert w >= 0.1 - 0.01
            assert w <= 0.6 + 0.01

    def test_too_few_assets_raises(self):
        with pytest.raises(ValueError):
            robust_optimize([0.1], [[0.04]], [0.02], 0.5)

    def test_mismatched_sigma_raises(self):
        with pytest.raises(ValueError):
            robust_optimize([0.1, 0.15], [[0.04]], [0.02, 0.03], 0.5)

    def test_mismatched_uncertainty_raises(self):
        mu, sigma, _ = _simple_inputs()
        with pytest.raises(ValueError):
            robust_optimize(mu, sigma, [0.02], 0.5)

    def test_non_finite_returns_raises(self):
        _, sigma, unc = _simple_inputs()
        with pytest.raises(ValueError):
            robust_optimize([float("nan"), 0.15], sigma, unc, 0.5)

    def test_non_finite_sigma_raises(self):
        mu, _, unc = _simple_inputs()
        with pytest.raises(ValueError):
            robust_optimize(mu, [[float("inf"), 0.0], [0.0, 0.09]], unc, 0.5)

    def test_negative_uncertainty_raises(self):
        mu, sigma, _ = _simple_inputs()
        with pytest.raises(ValueError):
            robust_optimize(mu, sigma, [-0.02, 0.03], 0.5)

    def test_negative_budget_raises(self):
        mu, sigma, unc = _simple_inputs()
        with pytest.raises(ValueError):
            robust_optimize(mu, sigma, unc, -0.5)

    def test_min_gt_max_raises(self):
        mu, sigma, unc = _simple_inputs()
        with pytest.raises(ValueError):
            robust_optimize(mu, sigma, unc, 0.5, max_weight=0.3, min_weight=0.8)

    def test_repr(self):
        mu, sigma, unc = _simple_inputs()
        result = robust_optimize(mu, sigma, unc, 0.5)
        r = repr(result)
        assert "RobustPortfolioResult" in r

    def test_higher_budget_lower_worst_case(self):
        mu, sigma, unc = _simple_inputs()
        r_low = robust_optimize(mu, sigma, unc, 0.1)
        r_high = robust_optimize(mu, sigma, unc, 2.0)
        assert r_high.worst_case_return <= r_low.worst_case_return + 1e-6

    def test_weights_count_matches_assets(self):
        mu, sigma, unc = _three_asset_inputs()
        result = robust_optimize(mu, sigma, unc, 0.5)
        assert len(result.weights) == len(mu)


# ===========================================================================
# worst_case_return
# ===========================================================================


class TestWorstCaseReturn:
    def test_basic(self):
        wc = worst_case_return([0.5, 0.5], [0.1, 0.2], [0.05, 0.05], 1.0)
        assert math.isfinite(wc)

    def test_no_uncertainty_equals_nominal(self):
        wc = worst_case_return([0.5, 0.5], [0.1, 0.2], [0.0, 0.0], 1.0)
        nominal = 0.5 * 0.1 + 0.5 * 0.2
        assert wc == pytest.approx(nominal, abs=1e-10)

    def test_with_uncertainty_less_than_nominal(self):
        wc = worst_case_return([0.5, 0.5], [0.1, 0.2], [0.05, 0.05], 1.0)
        nominal = 0.5 * 0.1 + 0.5 * 0.2
        assert wc < nominal + 1e-10

    def test_mismatched_lengths_raises(self):
        with pytest.raises(ValueError):
            worst_case_return([0.5, 0.5], [0.1], [0.05, 0.05], 1.0)

    def test_zero_budget_equals_nominal(self):
        wc = worst_case_return([0.5, 0.5], [0.1, 0.2], [0.05, 0.05], 0.0)
        nominal = 0.5 * 0.1 + 0.5 * 0.2
        assert wc == pytest.approx(nominal, abs=1e-10)

    def test_higher_budget_lower_return(self):
        wc1 = worst_case_return([0.5, 0.5], [0.1, 0.2], [0.05, 0.05], 0.5)
        wc2 = worst_case_return([0.5, 0.5], [0.1, 0.2], [0.05, 0.05], 2.0)
        assert wc2 <= wc1 + 1e-10

    def test_higher_uncertainty_lower_return(self):
        wc1 = worst_case_return([0.5, 0.5], [0.1, 0.2], [0.01, 0.01], 1.0)
        wc2 = worst_case_return([0.5, 0.5], [0.1, 0.2], [0.10, 0.10], 1.0)
        assert wc2 <= wc1 + 1e-10

    def test_three_assets(self):
        wc = worst_case_return(
            [0.3, 0.4, 0.3], [0.1, 0.15, 0.08], [0.02, 0.03, 0.01], 0.5
        )
        assert math.isfinite(wc)


# ===========================================================================
# robust_efficient_frontier
# ===========================================================================


class TestRobustEfficientFrontier:
    def test_basic(self):
        mu, sigma, unc = _simple_inputs()
        frontier = robust_efficient_frontier(mu, sigma, unc, 10)
        assert len(frontier) == 10

    def test_returns_tuples(self):
        mu, sigma, unc = _simple_inputs()
        frontier = robust_efficient_frontier(mu, sigma, unc, 5)
        for wc_ret, risk in frontier:
            assert math.isfinite(wc_ret)
            assert math.isfinite(risk)
            assert risk >= 0.0

    def test_too_few_points_raises(self):
        mu, sigma, unc = _simple_inputs()
        with pytest.raises(ValueError):
            robust_efficient_frontier(mu, sigma, unc, 1)

    def test_too_few_assets_raises(self):
        with pytest.raises(ValueError):
            robust_efficient_frontier([0.1], [[0.04]], [0.02], 5)

    def test_three_assets(self):
        mu, sigma, unc = _three_asset_inputs()
        frontier = robust_efficient_frontier(mu, sigma, unc, 10)
        assert len(frontier) == 10

    def test_many_points(self):
        mu, sigma, unc = _simple_inputs()
        frontier = robust_efficient_frontier(mu, sigma, unc, 50)
        assert len(frontier) == 50

    def test_risk_non_negative(self):
        mu, sigma, unc = _simple_inputs()
        frontier = robust_efficient_frontier(mu, sigma, unc, 10)
        for _, risk in frontier:
            assert risk >= -1e-10
