"""Tests for Almgren-Chriss optimal liquidation."""

import math
import pytest

from horizon._horizon import (
    ACTrajectory,
    ac_efficient_frontier,
    ac_kappa,
    ac_optimal_schedule,
)


# ===========================================================================
# ac_kappa
# ===========================================================================


class TestACKappa:
    def test_basic(self):
        k = ac_kappa(0.3, 0.01, 0.05, 1.0)
        assert k > 0.0
        assert math.isfinite(k)

    def test_higher_risk_aversion_higher_kappa(self):
        k1 = ac_kappa(0.3, 0.01, 0.05, 0.1)
        k2 = ac_kappa(0.3, 0.01, 0.05, 10.0)
        assert k2 > k1

    def test_higher_vol_higher_kappa(self):
        k1 = ac_kappa(0.1, 0.01, 0.05, 1.0)
        k2 = ac_kappa(0.5, 0.01, 0.05, 1.0)
        assert k2 > k1

    def test_higher_temp_impact_lower_kappa(self):
        k1 = ac_kappa(0.3, 0.01, 0.01, 1.0)
        k2 = ac_kappa(0.3, 0.01, 0.1, 1.0)
        assert k2 < k1

    def test_zero_vol_raises(self):
        with pytest.raises(ValueError):
            ac_kappa(0.0, 0.01, 0.05, 1.0)

    def test_negative_vol_raises(self):
        with pytest.raises(ValueError):
            ac_kappa(-0.3, 0.01, 0.05, 1.0)

    def test_negative_perm_impact_raises(self):
        with pytest.raises(ValueError):
            ac_kappa(0.3, -0.01, 0.05, 1.0)

    def test_zero_perm_impact_ok(self):
        k = ac_kappa(0.3, 0.0, 0.05, 1.0)
        assert k > 0.0

    def test_zero_temp_impact_raises(self):
        with pytest.raises(ValueError):
            ac_kappa(0.3, 0.01, 0.0, 1.0)

    def test_zero_risk_aversion_raises(self):
        with pytest.raises(ValueError):
            ac_kappa(0.3, 0.01, 0.05, 0.0)

    def test_nan_raises(self):
        with pytest.raises(ValueError):
            ac_kappa(float("nan"), 0.01, 0.05, 1.0)

    def test_inf_raises(self):
        with pytest.raises(ValueError):
            ac_kappa(float("inf"), 0.01, 0.05, 1.0)

    def test_negative_risk_aversion_raises(self):
        with pytest.raises(ValueError):
            ac_kappa(0.3, 0.01, 0.05, -1.0)

    def test_small_values(self):
        k = ac_kappa(0.01, 0.0001, 0.001, 0.01)
        assert k > 0.0
        assert math.isfinite(k)

    def test_large_values(self):
        k = ac_kappa(1.0, 0.1, 1.0, 10.0)
        assert k > 0.0
        assert math.isfinite(k)


# ===========================================================================
# ac_optimal_schedule
# ===========================================================================


class TestACOptimalSchedule:
    def test_basic(self):
        result = ac_optimal_schedule(1000.0, 10, 0.3, 0.01, 0.05, 1.0)
        assert isinstance(result, ACTrajectory)

    def test_result_fields(self):
        result = ac_optimal_schedule(1000.0, 10, 0.3, 0.01, 0.05, 1.0)
        assert hasattr(result, "schedule")
        assert hasattr(result, "expected_cost")
        assert hasattr(result, "variance")
        assert hasattr(result, "efficient_frontier")

    def test_schedule_length(self):
        result = ac_optimal_schedule(1000.0, 10, 0.3, 0.01, 0.05, 1.0)
        # Schedule has n_steps+1 entries (including initial)
        assert len(result.schedule) == 11

    def test_schedule_starts_at_total(self):
        result = ac_optimal_schedule(1000.0, 10, 0.3, 0.01, 0.05, 1.0)
        assert result.schedule[0] == pytest.approx(1000.0, rel=0.01)

    def test_schedule_ends_near_zero(self):
        result = ac_optimal_schedule(1000.0, 10, 0.3, 0.01, 0.05, 1.0)
        assert result.schedule[-1] == pytest.approx(0.0, abs=1.0)

    def test_schedule_monotonically_decreasing(self):
        result = ac_optimal_schedule(1000.0, 10, 0.3, 0.01, 0.05, 1.0)
        for i in range(len(result.schedule) - 1):
            assert result.schedule[i] >= result.schedule[i + 1] - 1e-6

    def test_expected_cost_non_negative(self):
        result = ac_optimal_schedule(1000.0, 10, 0.3, 0.01, 0.05, 1.0)
        assert result.expected_cost >= 0.0

    def test_variance_non_negative(self):
        result = ac_optimal_schedule(1000.0, 10, 0.3, 0.01, 0.05, 1.0)
        assert result.variance >= 0.0

    def test_efficient_frontier_exists(self):
        result = ac_optimal_schedule(1000.0, 10, 0.3, 0.01, 0.05, 1.0)
        assert len(result.efficient_frontier) > 0

    def test_zero_total_shares_raises(self):
        with pytest.raises(ValueError):
            ac_optimal_schedule(0.0, 10, 0.3, 0.01, 0.05, 1.0)

    def test_negative_total_shares_raises(self):
        with pytest.raises(ValueError):
            ac_optimal_schedule(-1000.0, 10, 0.3, 0.01, 0.05, 1.0)

    def test_zero_n_steps_raises(self):
        with pytest.raises(ValueError):
            ac_optimal_schedule(1000.0, 0, 0.3, 0.01, 0.05, 1.0)

    def test_zero_vol_raises(self):
        with pytest.raises(ValueError):
            ac_optimal_schedule(1000.0, 10, 0.0, 0.01, 0.05, 1.0)

    def test_zero_temp_impact_raises(self):
        with pytest.raises(ValueError):
            ac_optimal_schedule(1000.0, 10, 0.3, 0.01, 0.0, 1.0)

    def test_zero_risk_aversion_raises(self):
        with pytest.raises(ValueError):
            ac_optimal_schedule(1000.0, 10, 0.3, 0.01, 0.05, 0.0)

    def test_higher_risk_aversion_front_loaded(self):
        low_ra = ac_optimal_schedule(1000.0, 10, 0.3, 0.01, 0.05, 0.1)
        high_ra = ac_optimal_schedule(1000.0, 10, 0.3, 0.01, 0.05, 10.0)
        # Higher risk aversion -> more aggressive early trading
        assert high_ra.schedule[1] <= low_ra.schedule[1] + 1e-6

    def test_different_steps_finite_cost(self):
        fast = ac_optimal_schedule(1000.0, 5, 0.3, 0.01, 0.05, 1.0)
        slow = ac_optimal_schedule(1000.0, 50, 0.3, 0.01, 0.05, 1.0)
        # Both should produce finite positive costs
        assert fast.expected_cost > 0
        assert slow.expected_cost > 0

    def test_repr(self):
        result = ac_optimal_schedule(1000.0, 10, 0.3, 0.01, 0.05, 1.0)
        r = repr(result)
        assert "ACTrajectory" in r

    def test_nan_raises(self):
        with pytest.raises(ValueError):
            ac_optimal_schedule(float("nan"), 10, 0.3, 0.01, 0.05, 1.0)


# ===========================================================================
# ac_efficient_frontier
# ===========================================================================


class TestACEfficientFrontier:
    def test_basic(self):
        frontier = ac_efficient_frontier(1000.0, 10, 0.3, 0.01, 0.05, 5)
        assert len(frontier) == 5

    def test_returns_tuples(self):
        frontier = ac_efficient_frontier(1000.0, 10, 0.3, 0.01, 0.05, 5)
        for cost, var in frontier:
            assert isinstance(cost, float)
            assert isinstance(var, float)
            assert math.isfinite(cost)
            assert math.isfinite(var)

    def test_cost_non_negative(self):
        frontier = ac_efficient_frontier(1000.0, 10, 0.3, 0.01, 0.05, 10)
        for cost, var in frontier:
            assert cost >= 0.0

    def test_variance_non_negative(self):
        frontier = ac_efficient_frontier(1000.0, 10, 0.3, 0.01, 0.05, 10)
        for cost, var in frontier:
            assert var >= 0.0

    def test_zero_n_points_raises(self):
        with pytest.raises(ValueError):
            ac_efficient_frontier(1000.0, 10, 0.3, 0.01, 0.05, 0)

    def test_zero_total_shares_raises(self):
        with pytest.raises(ValueError):
            ac_efficient_frontier(0.0, 10, 0.3, 0.01, 0.05, 5)

    def test_zero_n_steps_raises(self):
        with pytest.raises(ValueError):
            ac_efficient_frontier(1000.0, 0, 0.3, 0.01, 0.05, 5)

    def test_zero_vol_raises(self):
        with pytest.raises(ValueError):
            ac_efficient_frontier(1000.0, 10, 0.0, 0.01, 0.05, 5)

    def test_sorted_by_cost(self):
        frontier = ac_efficient_frontier(1000.0, 10, 0.3, 0.01, 0.05, 10)
        costs = [c for c, _ in frontier]
        for i in range(len(costs) - 1):
            assert costs[i] <= costs[i + 1] + 1e-10

    def test_single_point(self):
        frontier = ac_efficient_frontier(1000.0, 10, 0.3, 0.01, 0.05, 1)
        assert len(frontier) == 1

    def test_many_points(self):
        frontier = ac_efficient_frontier(1000.0, 10, 0.3, 0.01, 0.05, 50)
        assert len(frontier) == 50
