"""Unit tests for v3 correlation-aware clustering.

Tests for compute_n_eff, compute_cluster_llr, and cluster_evidence functions.

Requirements validated: 4.1–4.5
"""
from __future__ import annotations

from datetime import datetime, timezone

import pytest

from services.aggregation.scoring import EvidenceUnit
from services.aggregation.worker import (
    cluster_evidence,
    compute_cluster_llr,
    compute_n_eff,
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_unit(cluster_id: str = "cluster_a", symbol: str = "AAPL") -> EvidenceUnit:
    """Create a minimal EvidenceUnit for testing."""
    return EvidenceUnit(
        symbol=symbol,
        layer="company",
        event_type="earnings",
        source_id="doc_1",
        source_group="reuters",
        timestamp=datetime(2024, 1, 15, 12, 0, tzinfo=timezone.utc),
        horizon="7d",
        direction=1,
        sentiment_strength=0.8,
        impact=0.7,
        extraction_conf=0.9,
        source_cred=0.85,
        novelty=0.6,
        event_base_rate=0.25,
        cluster_id=cluster_id,
    )


# ---------------------------------------------------------------------------
# Test: 3 identical articles from same source → n_eff < 3
# Requirement: 4.2, 4.3
# ---------------------------------------------------------------------------


class TestNEffIdenticalArticles:
    """3 identical articles from same wire (default rho=0.80) → n_eff < 3."""

    def test_n_eff_less_than_cluster_size(self):
        llrs = [1.0, 1.0, 1.0]
        # Default correlations: rho=0.80 for all pairs (same wire/source)
        n_eff = compute_n_eff(llrs)

        # Formula: (3)² / (3 + 2×3×0.80×1×1) = 9 / (3 + 4.8) = 9/7.8 ≈ 1.154
        expected = 9.0 / 7.8
        assert n_eff < 3.0
        assert n_eff == pytest.approx(expected, rel=1e-6)

    def test_n_eff_discounts_correlated_signals(self):
        """Higher correlation → lower n_eff."""
        llrs = [1.0, 1.0, 1.0]
        n_eff_correlated = compute_n_eff(llrs)  # default rho=0.80

        # Independent: rho=0.0
        identity = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
        n_eff_independent = compute_n_eff(llrs, correlations=identity)

        assert n_eff_correlated < n_eff_independent


# ---------------------------------------------------------------------------
# Test: 3 independent articles → n_eff ≈ 3
# Requirement: 4.2, 4.3
# ---------------------------------------------------------------------------


class TestNEffIndependentArticles:
    """3 independent articles (rho=0.0) → n_eff = 3.0."""

    def test_n_eff_equals_cluster_size(self):
        llrs = [1.0, 1.0, 1.0]
        # Zero off-diagonal correlations
        correlations = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
        n_eff = compute_n_eff(llrs, correlations=correlations)

        # n_eff = (3)² / (3 + 0) = 3.0
        assert n_eff == pytest.approx(3.0, rel=1e-6)

    def test_n_eff_with_varying_magnitudes(self):
        """Independent signals with different magnitudes still give n <= cluster size."""
        llrs = [0.5, 1.0, 2.0]
        correlations = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
        n_eff = compute_n_eff(llrs, correlations=correlations)

        # With zero correlations, n_eff = (sum |w|)^2 / sum(w^2)
        # = (0.5+1.0+2.0)^2 / (0.25+1.0+4.0) = 12.25 / 5.25 ≈ 2.333
        expected = (3.5**2) / (0.25 + 1.0 + 4.0)
        assert n_eff == pytest.approx(expected, rel=1e-6)
        assert n_eff <= 3.0


# ---------------------------------------------------------------------------
# Test: Single signal cluster → n_eff = 1.0
# Requirement: 4.2
# ---------------------------------------------------------------------------


class TestNEffSingleSignal:
    """Single signal in a cluster → n_eff = 1.0."""

    def test_single_signal(self):
        llrs = [0.5]
        n_eff = compute_n_eff(llrs)
        assert n_eff == 1.0

    def test_empty_cluster(self):
        llrs: list[float] = []
        n_eff = compute_n_eff(llrs)
        assert n_eff == 1.0


# ---------------------------------------------------------------------------
# Test: Cluster LLR clamp at ±2.5
# Requirement: 4.4, 4.5
# ---------------------------------------------------------------------------


class TestClusterLLRClamp:
    """Cluster LLR is clamped to [-2.5, 2.5]."""

    def test_positive_clamp(self):
        """Very large positive LLRs with high n_eff → clamped to 2.5."""
        llrs = [2.0, 2.0, 2.0, 2.0, 2.0]
        # Use independent correlations for max n_eff
        correlations = [
            [1, 0, 0, 0, 0],
            [0, 1, 0, 0, 0],
            [0, 0, 1, 0, 0],
            [0, 0, 0, 1, 0],
            [0, 0, 0, 0, 1],
        ]
        n_eff = compute_n_eff(llrs, correlations=correlations)
        cluster_llr = compute_cluster_llr(llrs, n_eff)

        # weighted_mean = 2.0, sqrt(5) ≈ 2.236, raw = 4.47 → clamp to 2.5
        assert cluster_llr == pytest.approx(2.5, rel=1e-6)

    def test_negative_clamp(self):
        """Very negative LLRs with high n_eff → clamped to -2.5."""
        llrs = [-2.0, -2.0, -2.0, -2.0, -2.0]
        correlations = [
            [1, 0, 0, 0, 0],
            [0, 1, 0, 0, 0],
            [0, 0, 1, 0, 0],
            [0, 0, 0, 1, 0],
            [0, 0, 0, 0, 1],
        ]
        n_eff = compute_n_eff(llrs, correlations=correlations)
        cluster_llr = compute_cluster_llr(llrs, n_eff)

        assert cluster_llr == pytest.approx(-2.5, rel=1e-6)

    def test_within_bounds_no_clamp(self):
        """Small LLRs with low n_eff → no clamping needed."""
        llrs = [0.3, 0.4]
        n_eff = compute_n_eff(llrs)
        cluster_llr = compute_cluster_llr(llrs, n_eff)

        assert -2.5 <= cluster_llr <= 2.5
        # Should NOT be at the clamp boundary
        assert abs(cluster_llr) < 2.5

    def test_single_signal_clamp(self):
        """Single signal beyond clamp → clamped."""
        llrs = [3.0]
        cluster_llr = compute_cluster_llr(llrs, n_eff=1.0)
        assert cluster_llr == pytest.approx(2.5, rel=1e-6)

    def test_single_signal_negative_clamp(self):
        """Single negative signal beyond clamp → clamped to -2.5."""
        llrs = [-3.0]
        cluster_llr = compute_cluster_llr(llrs, n_eff=1.0)
        assert cluster_llr == pytest.approx(-2.5, rel=1e-6)


# ---------------------------------------------------------------------------
# Test: All-zero LLRs → cluster_llr = 0.0
# Requirement: 4.4
# ---------------------------------------------------------------------------


class TestClusterLLRZero:
    """All-zero LLRs produce zero cluster LLR."""

    def test_all_zeros(self):
        llrs = [0.0, 0.0, 0.0]
        n_eff = compute_n_eff(llrs)
        cluster_llr = compute_cluster_llr(llrs, n_eff)
        assert cluster_llr == 0.0

    def test_empty_llrs(self):
        """Empty LLR list → 0.0."""
        cluster_llr = compute_cluster_llr([], n_eff=1.0)
        assert cluster_llr == 0.0


# ---------------------------------------------------------------------------
# Test: Grouping by correct key dimensions (cluster_id)
# Requirement: 4.1
# ---------------------------------------------------------------------------


class TestClusterEvidence:
    """cluster_evidence groups EvidenceUnits by cluster_id."""

    def test_grouping_by_cluster_id(self):
        """Units with same cluster_id are grouped together."""
        unit_a1 = _make_unit(cluster_id="cluster_a")
        unit_a2 = _make_unit(cluster_id="cluster_a")
        unit_b1 = _make_unit(cluster_id="cluster_b")

        units = [unit_a1, unit_a2, unit_b1]
        llrs = [1.0, 0.5, -0.3]

        clusters = cluster_evidence(units, llrs)

        assert len(clusters) == 2

        # Find clusters by id
        cluster_map = {c.cluster_id: c for c in clusters}

        assert "cluster_a" in cluster_map
        assert "cluster_b" in cluster_map

        # Cluster A has 2 units
        assert len(cluster_map["cluster_a"].units) == 2
        assert cluster_map["cluster_a"].llrs == [1.0, 0.5]

        # Cluster B has 1 unit
        assert len(cluster_map["cluster_b"].units) == 1
        assert cluster_map["cluster_b"].llrs == [-0.3]

    def test_single_cluster(self):
        """All units with same cluster_id → one cluster."""
        units = [_make_unit(cluster_id="only") for _ in range(4)]
        llrs = [0.1, 0.2, 0.3, 0.4]

        clusters = cluster_evidence(units, llrs)

        assert len(clusters) == 1
        assert clusters[0].cluster_id == "only"
        assert len(clusters[0].units) == 4
        assert clusters[0].llrs == [0.1, 0.2, 0.3, 0.4]

    def test_each_unit_different_cluster(self):
        """Each unit in its own cluster → N clusters."""
        units = [_make_unit(cluster_id=f"c_{i}") for i in range(5)]
        llrs = [0.1 * i for i in range(5)]

        clusters = cluster_evidence(units, llrs)

        assert len(clusters) == 5
        for c in clusters:
            assert len(c.units) == 1

    def test_empty_input(self):
        """No units → no clusters."""
        clusters = cluster_evidence([], [])
        assert clusters == []

    def test_llrs_parallel_to_units(self):
        """LLRs are correctly associated with their units."""
        unit_x = _make_unit(cluster_id="x")
        unit_y = _make_unit(cluster_id="y")
        unit_x2 = _make_unit(cluster_id="x")

        units = [unit_x, unit_y, unit_x2]
        llrs = [1.5, -0.7, 2.3]

        clusters = cluster_evidence(units, llrs)
        cluster_map = {c.cluster_id: c for c in clusters}

        assert cluster_map["x"].llrs == [1.5, 2.3]
        assert cluster_map["y"].llrs == [-0.7]