"""
Tests for microstructure/hawkes.py — multivariate Hawkes process calibrator.
"""
import numpy as np
from microstructure.hawkes import (
    HawkesCalibrator,
    hawkes_log_likelihood,
    hawkes_intensity,
    generate_hawkes_events,
)


class TestHawkesCalibrator:

    def test_initial_state(self):
        cal = HawkesCalibrator(n_dimensions=3)
        assert cal.n_dim == 3
        assert cal.mu.shape == (3,)
        assert cal.alpha.shape == (3, 3)
        assert cal.beta > 0

    def test_calibrate_on_synthetic_data(self):
        """Calibrate on synthetic events from known parameters, check recovery."""
        np.random.seed(42)
        # Generate events with known parameters: 3 types
        true_mu = np.array([0.5, 0.3, 0.2])
        true_alpha = np.array([
            [0.1, 0.05, 0.02],
            [0.03, 0.08, 0.01],
            [0.01, 0.02, 0.06],
        ])
        true_beta = 2.0
        max_time = 500.0

        events = generate_hawkes_events(true_mu, true_alpha, true_beta, max_time, seed=42)
        assert len(events) >= 3, f"Expected events across 3 types, got {len(events)}"

        cal = HawkesCalibrator(n_dimensions=3)
        cal.calibrate(events, max_time)

        # Check that calibrated parameters are within reasonable range
        assert np.all(cal.mu > 0), f"mu should be positive, got {cal.mu}"
        assert np.all(cal.alpha >= 0), f"alpha should be non-negative, got {cal.alpha}"
        assert cal.beta > 0

    def test_intensity_interpolation(self):
        """Intensity should recover to baseline between events and spike after."""
        cal = HawkesCalibrator(n_dimensions=3)
        cal.mu = np.array([0.5, 0.3, 0.2])
        cal.alpha = np.array([[0.1, 0, 0], [0, 0, 0], [0, 0, 0]])
        cal.beta = 2.0

        # Before first event at t=0, intensity = mu (baseline)
        i0 = cal.intensity(0, event_history=[], current_time=0.0)
        assert abs(i0 - cal.mu[0]) < 0.001

        # After an event of type 0 at t=0, type 0 intensity should spike
        i_after = cal.intensity(0, event_history=[(0, 0.0)], current_time=0.01)
        assert i_after > cal.mu[0]

        # After decay, should approach baseline
        i_later = cal.intensity(0, event_history=[(0, 0.0)], current_time=5.0)
        assert abs(i_later - cal.mu[0]) < 0.05

    def test_branching_ratio(self):
        """Branching ratio should be between 0 and 1."""
        cal = HawkesCalibrator(n_dimensions=3)
        cal.mu = np.array([0.5, 0.3, 0.2])
        cal.alpha = np.array([[0.1, 0, 0], [0, 0.1, 0], [0, 0, 0.1]])
        cal.beta = 2.0
        ratio = cal.branching_ratio()
        assert 0.0 <= ratio <= 1.0

    def test_forecast_activity(self):
        """Forecast event count in next window."""
        cal = HawkesCalibrator(n_dimensions=3)
        cal.mu = np.array([1.0, 0.5, 0.3])
        cal.alpha = np.array([[0.1, 0, 0], [0, 0, 0], [0, 0, 0]])
        cal.beta = 2.0

        events = [(0, 0.0), (0, 0.5), (0, 1.0), (1, 1.5)]
        forecast = cal.forecast_activity(events, current_time=2.0, horizon=5.0)
        assert len(forecast) == 3
        assert np.all(forecast >= 0)

    def test_cross_excitation_detected(self):
        """Alpha matrix should capture cross-excitation between types."""
        np.random.seed(123)
        true_mu = np.array([0.5, 0.3, 0.2])
        true_alpha = np.array([
            [0.2, 0.0, 0.0],
            [0.1, 0.1, 0.0],   # type 0 excites type 1
            [0.0, 0.0, 0.1],
        ])
        true_beta = 3.0

        events = generate_hawkes_events(true_mu, true_alpha, true_beta, 500.0, seed=123)
        cal = HawkesCalibrator(n_dimensions=3)
        cal.calibrate(events, 500.0)

        # Type 0 should excite type 1: alpha[1,0] > 0
        assert cal.alpha[1, 0] >= 0, f"Expected cross-excitation, got alpha[1,0]={cal.alpha[1,0]}"


class TestHawkesFunctions:

    def test_log_likelihood_improves_with_fit(self):
        np.random.seed(99)
        true_mu = np.array([0.5, 0.3, 0.2])
        true_alpha = np.array([[0.15, 0.0, 0.0], [0.0, 0.1, 0.0], [0.0, 0.0, 0.05]])
        true_beta = 2.5

        events = generate_hawkes_events(true_mu, true_alpha, true_beta, 200.0, seed=99)

        # Bad parameters
        bad_mu = np.array([0.1, 0.1, 0.1])
        bad_alpha = np.zeros((3, 3))
        ll_bad = hawkes_log_likelihood(events, bad_mu, bad_alpha, true_beta, 200.0)

        # Good parameters (close to truth)
        ll_good = hawkes_log_likelihood(events, true_mu, true_alpha, true_beta, 200.0)

        assert ll_good > ll_bad, f"Good params should give higher likelihood: {ll_good} vs {ll_bad}"

    def test_intensity_function(self):
        mu = np.array([0.5, 0.3])
        alpha = np.array([[0.1, 0.0], [0.0, 0.05]])
        beta = 2.0
        events = [(0, 0.0), (0, 0.5), (1, 1.0)]

        intensity = hawkes_intensity(mu, alpha, beta, events, current_time=0.6, dim=0)
        assert intensity > mu[0], f"After events, intensity should exceed baseline"

    def test_generate_events_produces_timestamps(self):
        np.random.seed(42)
        events = generate_hawkes_events(
            np.array([0.5, 0.3]),
            np.array([[0.1, 0.0], [0.0, 0.05]]),
            beta=2.0,
            max_time=100.0,
            seed=42,
        )
        assert len(events) > 0
        # Each event should be (type, time)
        for ev in events:
            assert isinstance(ev, tuple)
            assert len(ev) == 2
            assert ev[0] in (0, 1)
            assert 0 <= ev[1] <= 100.0