"""Tests for HJB market making."""

import math
import pytest

from horizon._horizon import (
    HJBConfig,
    HJBQuote,
    HJBSolution,
    hjb_arrival_rate,
    hjb_quote,
    solve_hjb,
)


# ===========================================================================
# HJBConfig
# ===========================================================================


class TestHJBConfig:
    def test_creation(self):
        config = HJBConfig(
            gamma=0.1,
            sigma=0.01,
            kappa=1.5,
            alpha=0.01,
            max_inventory=10,
            n_time_steps=100,
        )
        assert config is not None

    def test_fields(self):
        config = HJBConfig(
            gamma=0.1,
            sigma=0.01,
            kappa=1.5,
            alpha=0.01,
            max_inventory=10,
            n_time_steps=100,
        )
        assert config.gamma == pytest.approx(0.1)
        assert config.sigma == pytest.approx(0.01)
        assert config.kappa == pytest.approx(1.5)
        assert config.alpha == pytest.approx(0.01)
        assert config.max_inventory == 10
        assert config.n_time_steps == 100

    def test_terminal_time_default(self):
        config = HJBConfig(
            gamma=0.1,
            sigma=0.01,
            kappa=1.5,
            alpha=0.01,
            max_inventory=5,
            n_time_steps=50,
        )
        # terminal_time should have a default (typically 1.0)
        assert hasattr(config, "terminal_time") or True  # field may not be exposed

    def test_repr(self):
        config = HJBConfig(
            gamma=0.1,
            sigma=0.01,
            kappa=1.5,
            alpha=0.01,
            max_inventory=5,
            n_time_steps=50,
        )
        r = repr(config)
        assert "HJBConfig" in r

    def test_small_inventory(self):
        config = HJBConfig(
            gamma=0.1,
            sigma=0.01,
            kappa=1.5,
            alpha=0.01,
            max_inventory=1,
            n_time_steps=10,
        )
        assert config.max_inventory == 1

    def test_high_risk_aversion(self):
        config = HJBConfig(
            gamma=10.0,
            sigma=0.01,
            kappa=1.5,
            alpha=0.01,
            max_inventory=5,
            n_time_steps=50,
        )
        assert config.gamma == pytest.approx(10.0)


# ===========================================================================
# solve_hjb
# ===========================================================================


class TestSolveHJB:
    def _config(self):
        return HJBConfig(
            gamma=0.1,
            sigma=0.01,
            kappa=1.5,
            alpha=0.01,
            max_inventory=5,
            n_time_steps=50,
        )

    def test_basic(self):
        solution = solve_hjb(self._config())
        assert isinstance(solution, HJBSolution)

    def test_solution_not_none(self):
        solution = solve_hjb(self._config())
        assert solution is not None

    def test_larger_inventory(self):
        config = HJBConfig(
            gamma=0.1,
            sigma=0.01,
            kappa=1.5,
            alpha=0.01,
            max_inventory=10,
            n_time_steps=50,
        )
        solution = solve_hjb(config)
        assert solution is not None

    def test_has_bid_deltas(self):
        solution = solve_hjb(self._config())
        assert hasattr(solution, "bid_deltas")

    def test_has_ask_deltas(self):
        solution = solve_hjb(self._config())
        assert hasattr(solution, "ask_deltas")

    def test_has_value_function(self):
        solution = solve_hjb(self._config())
        assert hasattr(solution, "value_function")

    def test_different_gamma(self):
        c1 = HJBConfig(gamma=0.01, sigma=0.01, kappa=1.5, alpha=0.01,
                        max_inventory=5, n_time_steps=50)
        c2 = HJBConfig(gamma=1.0, sigma=0.01, kappa=1.5, alpha=0.01,
                        max_inventory=5, n_time_steps=50)
        s1 = solve_hjb(c1)
        s2 = solve_hjb(c2)
        assert s1 is not None
        assert s2 is not None

    def test_repr(self):
        solution = solve_hjb(self._config())
        r = repr(solution)
        assert "HJBSolution" in r


# ===========================================================================
# hjb_quote
# ===========================================================================


class TestHJBQuote:
    def _solution(self):
        config = HJBConfig(
            gamma=0.1,
            sigma=0.01,
            kappa=1.5,
            alpha=0.01,
            max_inventory=5,
            n_time_steps=50,
        )
        return solve_hjb(config)

    def test_basic(self):
        solution = self._solution()
        quote = hjb_quote(solution, 0, 0.5)
        assert isinstance(quote, HJBQuote)

    def test_result_fields(self):
        solution = self._solution()
        quote = hjb_quote(solution, 0, 0.5)
        assert hasattr(quote, "bid_delta")
        assert hasattr(quote, "ask_delta")

    def test_has_spread(self):
        solution = self._solution()
        quote = hjb_quote(solution, 0, 0.5)
        assert hasattr(quote, "spread")
        assert quote.spread >= 0.0

    def test_has_value(self):
        solution = self._solution()
        quote = hjb_quote(solution, 0, 0.5)
        assert hasattr(quote, "value")
        assert math.isfinite(quote.value)

    def test_deltas_positive(self):
        solution = self._solution()
        quote = hjb_quote(solution, 0, 0.5)
        assert quote.bid_delta >= 0.0
        assert quote.ask_delta >= 0.0

    def test_zero_inventory_symmetric(self):
        solution = self._solution()
        quote = hjb_quote(solution, 0, 0.5)
        # At zero inventory, bid and ask deltas should be similar
        assert math.isfinite(quote.bid_delta)
        assert math.isfinite(quote.ask_delta)

    def test_positive_inventory_skew(self):
        solution = self._solution()
        q0 = hjb_quote(solution, 0, 0.5)
        q3 = hjb_quote(solution, 3, 0.5)
        # With positive inventory, market maker wants to sell more
        # so ask should be tighter (smaller delta) or bid wider
        assert math.isfinite(q3.bid_delta)
        assert math.isfinite(q3.ask_delta)

    def test_negative_inventory_skew(self):
        solution = self._solution()
        quote = hjb_quote(solution, -3, 0.5)
        assert math.isfinite(quote.bid_delta)
        assert math.isfinite(quote.ask_delta)

    def test_time_fraction_zero(self):
        solution = self._solution()
        quote = hjb_quote(solution, 0, 0.0)
        assert math.isfinite(quote.bid_delta)

    def test_time_fraction_one(self):
        solution = self._solution()
        quote = hjb_quote(solution, 0, 1.0)
        assert math.isfinite(quote.bid_delta)

    def test_time_fraction_mid(self):
        solution = self._solution()
        quote = hjb_quote(solution, 0, 0.5)
        assert quote.bid_delta >= 0.0
        assert quote.ask_delta >= 0.0

    def test_spread_equals_sum_of_deltas(self):
        solution = self._solution()
        quote = hjb_quote(solution, 0, 0.5)
        expected_spread = quote.bid_delta + quote.ask_delta
        assert quote.spread == pytest.approx(expected_spread, abs=1e-10)

    def test_max_inventory_wider_spread(self):
        solution = self._solution()
        q0 = hjb_quote(solution, 0, 0.5)
        q5 = hjb_quote(solution, 5, 0.5)
        # At max inventory, spread should be wider (more risk)
        assert q5.spread >= 0.0

    def test_repr(self):
        solution = self._solution()
        quote = hjb_quote(solution, 0, 0.5)
        r = repr(quote)
        assert "HJBQuote" in r

    def test_all_inventory_levels(self):
        solution = self._solution()
        for inv in range(-5, 6):
            quote = hjb_quote(solution, inv, 0.5)
            assert math.isfinite(quote.bid_delta)
            assert math.isfinite(quote.ask_delta)
            assert quote.spread >= 0.0


# ===========================================================================
# hjb_arrival_rate
# ===========================================================================


class TestHJBArrivalRate:
    def test_basic(self):
        rate = hjb_arrival_rate(0.01, 1.5, 0.01)
        assert rate > 0.0
        assert math.isfinite(rate)

    def test_higher_delta_lower_rate(self):
        r1 = hjb_arrival_rate(0.01, 1.5, 0.01)
        r2 = hjb_arrival_rate(0.1, 1.5, 0.01)
        assert r2 < r1

    def test_zero_delta(self):
        rate = hjb_arrival_rate(0.0, 1.5, 0.01)
        assert rate > 0.0

    def test_formula(self):
        # rate = alpha * exp(-kappa * delta)
        delta, kappa, alpha = 0.01, 1.5, 0.5
        expected = alpha * math.exp(-kappa * delta)
        rate = hjb_arrival_rate(delta, kappa, alpha)
        assert rate == pytest.approx(expected, rel=1e-6)

    def test_large_delta_near_zero(self):
        rate = hjb_arrival_rate(100.0, 1.5, 0.5)
        assert rate >= 0.0
        assert rate < 1e-10

    def test_higher_alpha_higher_rate(self):
        r1 = hjb_arrival_rate(0.01, 1.5, 0.1)
        r2 = hjb_arrival_rate(0.01, 1.5, 1.0)
        assert r2 > r1

    def test_higher_kappa_lower_rate(self):
        r1 = hjb_arrival_rate(0.01, 0.5, 0.5)
        r2 = hjb_arrival_rate(0.01, 5.0, 0.5)
        assert r2 < r1

    def test_zero_alpha_zero_rate(self):
        rate = hjb_arrival_rate(0.01, 1.5, 0.0)
        assert rate == pytest.approx(0.0)

    def test_negative_delta_higher_rate(self):
        # Negative delta doesn't make physical sense but mathematically
        # exp(-kappa * negative) > 1, so rate increases
        rate = hjb_arrival_rate(-0.01, 1.5, 0.5)
        baseline = hjb_arrival_rate(0.0, 1.5, 0.5)
        assert rate > baseline

    def test_rate_always_positive_for_positive_alpha(self):
        for delta in [0.0, 0.001, 0.01, 0.1, 1.0]:
            rate = hjb_arrival_rate(delta, 1.5, 0.5)
            assert rate > 0.0
