"""Tests for HawkesProcess, ledoit_wolf_shrinkage, and their pipeline wrappers."""

import math
import time

import pytest

from horizon._horizon import HawkesProcess, Market, ledoit_wolf_shrinkage
from horizon.context import Context, FeedData
from horizon.quant import correlation_estimator, hawkes_intensity


# ===========================================================================
# HawkesProcess -- Rust type
# ===========================================================================


class TestHawkesProcessCreation:
    def test_default_params(self):
        hp = HawkesProcess()
        assert hp.branching_ratio() >= 0.0
        assert hp.event_count() == 0

    def test_custom_params(self):
        hp = HawkesProcess(mu=0.2, alpha=0.8, beta=2.0, max_events=500)
        assert hp.event_count() == 0

    def test_repr(self):
        hp = HawkesProcess(mu=0.1, alpha=0.5, beta=1.0)
        r = repr(hp)
        assert "HawkesProcess" in r


class TestHawkesIntensity:
    def test_intensity_increases_after_events(self):
        hp = HawkesProcess(mu=0.1, alpha=0.5, beta=1.0)
        t = 1000.0
        baseline = hp.intensity(t)
        hp.add_event(t)
        after = hp.intensity(t + 0.001)
        assert after > baseline

    def test_intensity_decay(self):
        hp = HawkesProcess(mu=0.1, alpha=0.5, beta=1.0)
        t = 1000.0
        hp.add_event(t)
        soon = hp.intensity(t + 0.01)
        later = hp.intensity(t + 10.0)
        assert soon > later

    def test_intensity_now(self):
        hp = HawkesProcess(mu=0.1, alpha=0.5, beta=1.0)
        val = hp.intensity_now()
        # With no events, intensity should be approximately mu
        assert val >= 0.0
        assert math.isfinite(val)

    def test_multiple_events_compound(self):
        hp = HawkesProcess(mu=0.1, alpha=0.5, beta=1.0)
        t = 1000.0
        hp.add_event(t)
        one_event = hp.intensity(t + 0.001)
        hp.add_event(t + 0.0005)
        two_events = hp.intensity(t + 0.001)
        assert two_events > one_event


class TestHawkesBranchingRatio:
    def test_branching_ratio_value(self):
        hp = HawkesProcess(mu=0.1, alpha=0.5, beta=1.0)
        br = hp.branching_ratio()
        # alpha/beta = 0.5/1.0 = 0.5
        assert br == pytest.approx(0.5, abs=1e-10)

    def test_subcritical(self):
        hp = HawkesProcess(mu=0.1, alpha=0.3, beta=1.0)
        assert hp.branching_ratio() < 1.0

    def test_critical(self):
        hp = HawkesProcess(mu=0.1, alpha=1.0, beta=1.0)
        assert hp.branching_ratio() == pytest.approx(1.0, abs=1e-10)


class TestHawkesExpectedEvents:
    def test_expected_positive(self):
        hp = HawkesProcess(mu=0.1, alpha=0.5, beta=1.0)
        t = 1000.0
        hp.add_event(t)
        expected = hp.expected_events(t, 60.0)
        assert expected > 0.0
        assert math.isfinite(expected)

    def test_longer_horizon_more_events(self):
        hp = HawkesProcess(mu=0.1, alpha=0.5, beta=1.0)
        t = 1000.0
        hp.add_event(t)
        e_short = hp.expected_events(t, 10.0)
        e_long = hp.expected_events(t, 100.0)
        assert e_long >= e_short


class TestHawkesReset:
    def test_reset_clears_events(self):
        hp = HawkesProcess(mu=0.1, alpha=0.5, beta=1.0)
        t = 1000.0
        hp.add_event(t)
        hp.add_event(t + 1.0)
        assert hp.event_count() == 2
        hp.reset()
        assert hp.event_count() == 0

    def test_intensity_returns_to_baseline_after_reset(self):
        hp = HawkesProcess(mu=0.1, alpha=0.5, beta=1.0)
        t = 1000.0
        hp.add_event(t)
        hp.reset()
        intensity = hp.intensity(t + 0.001)
        # Should be back to approximately mu (no events)
        assert intensity == pytest.approx(0.1, abs=0.05)


# ===========================================================================
# HawkesProcess -- Pipeline wrapper
# ===========================================================================


class TestHawkesIntensityPipeline:
    def test_injects_params(self):
        fn = hawkes_intensity(feed_name="btc", mu=0.1, alpha=0.5, beta=1.0)
        ctx = Context(
            feeds={"btc": FeedData(price=0.5)},
            market=Market(id="mkt1"),
            params={},
        )
        fn(ctx)
        assert "hawkes_intensity" in ctx.params
        assert "hawkes_branching" in ctx.params
        assert "hawkes_expected_1m" in ctx.params
        assert isinstance(ctx.params["hawkes_intensity"], float)
        assert isinstance(ctx.params["hawkes_branching"], float)
        assert isinstance(ctx.params["hawkes_expected_1m"], float)

    def test_fill_triggers_event(self):
        fn = hawkes_intensity(feed_name="btc", mu=0.1, alpha=0.5, beta=1.0)
        # First call: no fill, establishes baseline
        ctx_no_fill = Context(
            feeds={"btc": FeedData(price=0.5)},
            market=Market(id="mkt1"),
            params={},
        )
        fn(ctx_no_fill)
        baseline = ctx_no_fill.params["hawkes_intensity"]

        # Second call: with fill, intensity should increase
        ctx_fill = Context(
            feeds={"btc": FeedData(price=0.5)},
            market=Market(id="mkt1"),
            params={"fills_this_cycle": 3},
        )
        fn(ctx_fill)
        after_fill = ctx_fill.params["hawkes_intensity"]
        # Events were added this tick and intensity is queried at now+epsilon
        assert after_fill > baseline

    def test_price_jump_triggers_event(self):
        fn = hawkes_intensity(feed_name="btc", mu=0.1, alpha=0.5, beta=1.0)
        # First tick establishes baseline price
        ctx1 = Context(
            feeds={"btc": FeedData(price=0.50)},
            market=Market(id="mkt1"),
            params={},
        )
        fn(ctx1)
        intensity_before = ctx1.params["hawkes_intensity"]

        # Second tick with >2% price jump (10%)
        ctx2 = Context(
            feeds={"btc": FeedData(price=0.55)},
            market=Market(id="mkt1"),
            params={},
        )
        fn(ctx2)
        intensity_after = ctx2.params["hawkes_intensity"]
        assert intensity_after > intensity_before

    def test_small_price_change_no_event(self):
        fn = hawkes_intensity(feed_name="btc", mu=0.1, alpha=0.5, beta=1.0)
        # First tick
        ctx1 = Context(
            feeds={"btc": FeedData(price=0.50)},
            market=Market(id="mkt1"),
            params={},
        )
        fn(ctx1)
        i1 = ctx1.params["hawkes_intensity"]

        # Second tick: 1% change (below 2% threshold)
        ctx2 = Context(
            feeds={"btc": FeedData(price=0.505)},
            market=Market(id="mkt1"),
            params={},
        )
        fn(ctx2)
        i2 = ctx2.params["hawkes_intensity"]
        # Intensity should not have increased beyond natural decay
        # (it should be approximately the same as the baseline mu)
        assert i2 <= i1 + 1e-6

    def test_no_feed_graceful(self):
        fn = hawkes_intensity(feed_name="btc")
        ctx = Context(
            feeds={},
            market=Market(id="mkt1"),
            params={},
        )
        fn(ctx)
        assert "hawkes_intensity" in ctx.params
        assert math.isfinite(ctx.params["hawkes_intensity"])

    def test_per_market_isolation(self):
        fn = hawkes_intensity(feed_name="btc", mu=0.1, alpha=0.5, beta=1.0)
        # Market 1: add fills
        ctx1 = Context(
            feeds={"btc": FeedData(price=0.5)},
            market=Market(id="mkt1"),
            params={"fills_this_cycle": 5},
        )
        fn(ctx1)
        intensity_mkt1 = ctx1.params["hawkes_intensity"]

        # Market 2: no fills
        ctx2 = Context(
            feeds={"btc": FeedData(price=0.5)},
            market=Market(id="mkt2"),
            params={},
        )
        fn(ctx2)
        intensity_mkt2 = ctx2.params["hawkes_intensity"]

        # Market 1 had fills, market 2 did not
        assert intensity_mkt1 > intensity_mkt2

    def test_custom_param_name(self):
        fn = hawkes_intensity(feed_name="btc", param_name="my_hawkes")
        ctx = Context(
            feeds={"btc": FeedData(price=0.5)},
            market=Market(id="mkt1"),
            params={},
        )
        fn(ctx)
        assert "my_hawkes_intensity" in ctx.params
        assert "my_hawkes_branching" in ctx.params
        assert "my_hawkes_expected_1m" in ctx.params

    def test_returns_none(self):
        fn = hawkes_intensity(feed_name="btc")
        ctx = Context(
            feeds={"btc": FeedData(price=0.5)},
            market=Market(id="mkt1"),
            params={},
        )
        result = fn(ctx)
        assert result is None


# ===========================================================================
# Ledoit-Wolf shrinkage -- Rust function
# ===========================================================================


class TestLedoitWolfShrinkage:
    def test_uncorrelated_data(self):
        # 100 observations of 2 uncorrelated assets
        data = []
        for i in range(100):
            data.append([
                math.sin(i * 0.1) * 0.01,
                math.cos(i * 0.7 + 3.0) * 0.01,
            ])
        matrix, shrinkage = ledoit_wolf_shrinkage(data)
        assert len(matrix) == 2
        assert len(matrix[0]) == 2
        # Off-diagonals should be small relative to diagonals
        assert abs(matrix[0][1]) < abs(matrix[0][0])
        assert abs(matrix[1][0]) < abs(matrix[1][1])

    def test_correlated_data(self):
        # 50 observations of 2 correlated assets
        data = [[0.01 * i, 0.01 * i + 0.001] for i in range(50)]
        matrix, shrinkage = ledoit_wolf_shrinkage(data)
        assert len(matrix) == 2
        assert abs(matrix[0][1]) > 0.0
        assert abs(matrix[1][0]) > 0.0

    def test_shrinkage_in_range(self):
        data = []
        for i in range(100):
            data.append([
                math.sin(i * 0.1) * 0.01,
                math.cos(i * 0.3) * 0.01,
            ])
        _, shrinkage = ledoit_wolf_shrinkage(data)
        assert 0.0 <= shrinkage <= 1.0

    def test_symmetric_matrix(self):
        # 50 observations of 3 assets
        data = []
        for i in range(50):
            data.append([
                math.sin(i * 0.1),
                math.cos(i * 0.2),
                math.sin(i * 0.3 + 1),
            ])
        matrix, _ = ledoit_wolf_shrinkage(data)
        assert len(matrix) == 3
        for i in range(3):
            for j in range(3):
                assert matrix[i][j] == pytest.approx(matrix[j][i], abs=1e-12)

    def test_positive_diagonal(self):
        data = []
        for i in range(100):
            data.append([
                math.sin(i * 0.1) * 0.01,
                math.cos(i * 0.3) * 0.01,
            ])
        matrix, _ = ledoit_wolf_shrinkage(data)
        assert matrix[0][0] > 0.0
        assert matrix[1][1] > 0.0

    def test_too_few_observations(self):
        # Only 1 observation of 2 assets -- should raise
        with pytest.raises((ValueError, RuntimeError)):
            ledoit_wolf_shrinkage([[0.01, 0.02]])

    def test_single_asset(self):
        # 50 observations of 1 asset (1x1 matrix)
        data = [[math.sin(i * 0.1) * 0.01] for i in range(50)]
        matrix, shrinkage = ledoit_wolf_shrinkage(data)
        assert len(matrix) == 1
        assert len(matrix[0]) == 1
        assert matrix[0][0] > 0.0

    def test_three_assets_dimensions(self):
        # 30 observations of 3 assets -> 3x3 matrix
        data = []
        for i in range(30):
            data.append([i * 0.01, i * 0.02, i * 0.005])
        matrix, shrinkage = ledoit_wolf_shrinkage(data)
        assert len(matrix) == 3
        assert all(len(row) == 3 for row in matrix)


# ===========================================================================
# correlation_estimator -- Pipeline wrapper
# ===========================================================================


class TestCorrelationEstimatorPipeline:
    def test_injects_params(self):
        fn = correlation_estimator(feed_names=["a", "b"], window=5)
        # Run enough ticks to trigger computation
        for i in range(6):
            ctx = Context(
                feeds={
                    "a": FeedData(price=0.5 + i * 0.01),
                    "b": FeedData(price=0.6 + i * 0.005),
                },
                market=Market(id="mkt1"),
                params={},
            )
            fn(ctx)
        assert "correlation_matrix" in ctx.params
        assert "correlation_shrinkage_intensity" in ctx.params
        assert "correlation_updated" in ctx.params

    def test_matrix_dimensions(self):
        fn = correlation_estimator(feed_names=["a", "b"], window=5)
        for i in range(6):
            ctx = Context(
                feeds={
                    "a": FeedData(price=0.5 + i * 0.01),
                    "b": FeedData(price=0.6 + i * 0.005),
                },
                market=Market(id="mkt1"),
                params={},
            )
            fn(ctx)
        matrix = ctx.params["correlation_matrix"]
        if matrix:  # If updated
            assert len(matrix) == 2
            assert len(matrix[0]) == 2

    def test_no_update_before_window(self):
        fn = correlation_estimator(feed_names=["a", "b"], window=100)
        ctx = Context(
            feeds={
                "a": FeedData(price=0.5),
                "b": FeedData(price=0.6),
            },
            market=Market(id="mkt1"),
            params={},
        )
        fn(ctx)
        assert ctx.params["correlation_updated"] is False

    def test_empty_feeds(self):
        fn = correlation_estimator(feed_names=["a", "b"], window=5)
        ctx = Context(feeds={}, market=Market(id="mkt1"), params={})
        fn(ctx)
        assert ctx.params["correlation_updated"] is False

    def test_custom_param_name(self):
        fn = correlation_estimator(feed_names=["a", "b"], window=5, param_name="corr")
        for i in range(6):
            ctx = Context(
                feeds={
                    "a": FeedData(price=0.5 + i * 0.01),
                    "b": FeedData(price=0.6 + i * 0.005),
                },
                market=Market(id="mkt1"),
                params={},
            )
            fn(ctx)
        assert "corr_matrix" in ctx.params
        assert "corr_shrinkage_intensity" in ctx.params
        assert "corr_updated" in ctx.params

    def test_auto_detect_feeds(self):
        fn = correlation_estimator(feed_names=None, window=5)
        for i in range(6):
            ctx = Context(
                feeds={
                    "x": FeedData(price=0.5 + i * 0.01),
                    "y": FeedData(price=0.6 + i * 0.005),
                },
                market=Market(id="mkt1"),
                params={},
            )
            fn(ctx)
        assert "correlation_matrix" in ctx.params

    def test_returns_none(self):
        fn = correlation_estimator(feed_names=["a", "b"], window=5)
        ctx = Context(
            feeds={
                "a": FeedData(price=0.5),
                "b": FeedData(price=0.6),
            },
            market=Market(id="mkt1"),
            params={},
        )
        result = fn(ctx)
        assert result is None
