"""Tests for cross-impact matrix estimation."""

import math
import pytest

from horizon._horizon import (
    CrossImpactEstimator,
    CrossImpactMatrix,
    estimate_cross_impact,
)


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

def _returns_volume(n_markets=3, n_obs=30, seed=42):
    """Generate pseudo-random returns and volumes.
    Returns (returns_matrix, volume_matrix) where each is n_markets x n_obs."""
    import random
    rng = random.Random(seed)
    returns = [[] for _ in range(n_markets)]
    volumes = [[] for _ in range(n_markets)]
    for _ in range(n_obs):
        for m in range(n_markets):
            returns[m].append(rng.gauss(0, 0.01))
            volumes[m].append(rng.random() * 100)
    return returns, volumes


# ===========================================================================
# estimate_cross_impact
# ===========================================================================


class TestEstimateCrossImpact:
    def test_basic(self):
        returns, volumes = _returns_volume()
        names = ["A", "B", "C"]
        result = estimate_cross_impact(returns, volumes, names)
        assert isinstance(result, CrossImpactMatrix)

    def test_result_fields(self):
        returns, volumes = _returns_volume()
        names = ["A", "B", "C"]
        result = estimate_cross_impact(returns, volumes, names)
        assert hasattr(result, "matrix")
        assert hasattr(result, "market_names")
        assert hasattr(result, "eigenvalues")
        assert hasattr(result, "is_psd")

    def test_matrix_shape(self):
        returns, volumes = _returns_volume(n_markets=3)
        names = ["A", "B", "C"]
        result = estimate_cross_impact(returns, volumes, names)
        assert len(result.matrix) == 3
        for row in result.matrix:
            assert len(row) == 3

    def test_market_names(self):
        returns, volumes = _returns_volume()
        names = ["X", "Y", "Z"]
        result = estimate_cross_impact(returns, volumes, names)
        assert result.market_names == names

    def test_eigenvalues_count(self):
        returns, volumes = _returns_volume(n_markets=3)
        names = ["A", "B", "C"]
        result = estimate_cross_impact(returns, volumes, names)
        assert len(result.eigenvalues) == 3

    def test_two_markets(self):
        returns, volumes = _returns_volume(n_markets=2)
        names = ["A", "B"]
        result = estimate_cross_impact(returns, volumes, names)
        assert len(result.matrix) == 2

    def test_four_markets(self):
        returns, volumes = _returns_volume(n_markets=4, n_obs=50)
        names = ["A", "B", "C", "D"]
        result = estimate_cross_impact(returns, volumes, names)
        assert len(result.matrix) == 4
        assert len(result.eigenvalues) == 4

    def test_matrix_entries_finite(self):
        returns, volumes = _returns_volume()
        names = ["A", "B", "C"]
        result = estimate_cross_impact(returns, volumes, names)
        for row in result.matrix:
            for val in row:
                assert math.isfinite(val)

    def test_eigenvalues_finite(self):
        returns, volumes = _returns_volume()
        names = ["A", "B", "C"]
        result = estimate_cross_impact(returns, volumes, names)
        for ev in result.eigenvalues:
            assert math.isfinite(ev)

    def test_is_psd_is_bool(self):
        returns, volumes = _returns_volume()
        names = ["A", "B", "C"]
        result = estimate_cross_impact(returns, volumes, names)
        assert isinstance(result.is_psd, bool)

    def test_mismatched_names_raises(self):
        returns, volumes = _returns_volume(n_markets=3)
        with pytest.raises(ValueError):
            estimate_cross_impact(returns, volumes, ["A", "B"])

    def test_mismatched_volumes_raises(self):
        returns, volumes = _returns_volume(n_markets=3)
        with pytest.raises(ValueError):
            estimate_cross_impact(returns, volumes[:2], ["A", "B", "C"])

    def test_repr(self):
        returns, volumes = _returns_volume()
        names = ["A", "B", "C"]
        result = estimate_cross_impact(returns, volumes, names)
        r = repr(result)
        assert "CrossImpactMatrix" in r


# ===========================================================================
# CrossImpactEstimator (streaming)
# ===========================================================================


class TestCrossImpactEstimator:
    def test_creation(self):
        est = CrossImpactEstimator(3, 20)
        assert est is not None

    def test_add_observation(self):
        est = CrossImpactEstimator(3, 20)
        est.add_observation([0.01, -0.02, 0.005], [100.0, 200.0, 150.0])

    def test_estimate_after_observations(self):
        est = CrossImpactEstimator(3, 20)
        import random
        rng = random.Random(42)
        for _ in range(25):
            ret = [rng.gauss(0, 0.01) for _ in range(3)]
            vol = [rng.random() * 100 for _ in range(3)]
            est.add_observation(ret, vol)
        result = est.estimate()
        assert isinstance(result, CrossImpactMatrix)

    def test_estimate_insufficient_data(self):
        est = CrossImpactEstimator(3, 20)
        est.add_observation([0.01, -0.02, 0.005], [100.0, 200.0, 150.0])
        # May return a result or raise - depends on implementation
        try:
            result = est.estimate()
        except (ValueError, RuntimeError):
            pass

    def test_wrong_dimension_raises(self):
        est = CrossImpactEstimator(3, 20)
        with pytest.raises(ValueError):
            est.add_observation([0.01, -0.02], [100.0, 200.0, 150.0])

    def test_repr(self):
        est = CrossImpactEstimator(3, 20)
        r = repr(est)
        assert "CrossImpactEstimator" in r

    def test_lookback_parameter(self):
        est = CrossImpactEstimator(2, 10)
        import random
        rng = random.Random(42)
        for _ in range(15):
            ret = [rng.gauss(0, 0.01) for _ in range(2)]
            vol = [rng.random() * 100 for _ in range(2)]
            est.add_observation(ret, vol)
        result = est.estimate()
        assert len(result.matrix) == 2

    def test_streaming_result_matches_batch(self):
        import random
        rng = random.Random(42)
        n_markets = 3
        n_obs = 30
        returns = [[] for _ in range(n_markets)]
        volumes = [[] for _ in range(n_markets)]
        est = CrossImpactEstimator(n_markets, n_obs)
        for _ in range(n_obs):
            ret = [rng.gauss(0, 0.01) for _ in range(n_markets)]
            vol = [rng.random() * 100 for _ in range(n_markets)]
            for m in range(n_markets):
                returns[m].append(ret[m])
                volumes[m].append(vol[m])
            est.add_observation(ret, vol)
        streaming = est.estimate()
        batch = estimate_cross_impact(returns, volumes, ["A", "B", "C"])
        # Both should produce matrices of same shape
        assert len(streaming.matrix) == len(batch.matrix)
