"""Unit tests for services.signal_engine.correlation — Signal cluster classification and penalty.

Tests classify_signal mapping, apply_correlation_penalty decay logic,
cross-cluster independence, single-signal clusters, and edge cases.

Requirements: 7.1, 7.2, 7.3, 7.4
"""
from __future__ import annotations

import math

from services.signal_engine.correlation import (
    SignalCluster,
    apply_correlation_penalty,
    classify_signal,
)
from services.signal_engine.models import LikelihoodRatio

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


def _lr(
    signal_type: str,
    cluster: str,
    log_lr: float,
    *,
    hit_rate: float = 0.6,
    strength: float = 0.7,
) -> LikelihoodRatio:
    """Create a LikelihoodRatio with sensible defaults."""
    return LikelihoodRatio(
        signal_type=signal_type,
        cluster=cluster,
        lr=math.exp(log_lr),
        log_lr=log_lr,
        penalized_log_lr=log_lr,  # pre-penalty: same as log_lr
        hit_rate=hit_rate,
        strength=strength,
    )


# ===========================================================================
# 1. classify_signal — known signal types (Requirement 7.1)
# ===========================================================================


class TestClassifySignal:
    """Verify signal type → cluster mapping."""

    def test_ma_stack_is_momentum(self) -> None:
        assert classify_signal("ma_stack") == SignalCluster.MOMENTUM

    def test_rsi_is_momentum(self) -> None:
        assert classify_signal("rsi") == SignalCluster.MOMENTUM

    def test_fibonacci_is_structure(self) -> None:
        assert classify_signal("fibonacci") == SignalCluster.STRUCTURE

    def test_elliott_wave_is_structure(self) -> None:
        assert classify_signal("elliott_wave") == SignalCluster.STRUCTURE

    def test_cup_handle_is_structure(self) -> None:
        assert classify_signal("cup_handle") == SignalCluster.STRUCTURE

    def test_atr_is_volatility(self) -> None:
        assert classify_signal("atr") == SignalCluster.VOLATILITY

    def test_bollinger_is_volatility(self) -> None:
        assert classify_signal("bollinger") == SignalCluster.VOLATILITY

    def test_valuation_is_fundamentals(self) -> None:
        assert classify_signal("valuation") == SignalCluster.FUNDAMENTALS

    def test_earnings_is_fundamentals(self) -> None:
        assert classify_signal("earnings") == SignalCluster.FUNDAMENTALS

    def test_macro_is_fundamentals(self) -> None:
        assert classify_signal("macro") == SignalCluster.FUNDAMENTALS

    def test_unknown_signal_defaults_to_fundamentals(self) -> None:
        """Unknown signal types fall back to FUNDAMENTALS."""
        assert classify_signal("unknown_xyz") == SignalCluster.FUNDAMENTALS


# ===========================================================================
# 2. apply_correlation_penalty — within-cluster decay (Requirement 7.2)
# ===========================================================================


class TestWithinClusterDecay:
    """Within a cluster, strongest LR at full weight, subsequent at 0.5^(n-1)."""

    def test_two_momentum_signals_decay(self) -> None:
        """Second signal in same cluster gets 0.5 decay."""
        lrs = [
            _lr("ma_stack", "momentum", log_lr=0.8),
            _lr("rsi", "momentum", log_lr=0.5),
        ]
        result = apply_correlation_penalty(lrs)

        # ma_stack is strongest (0.8 > 0.5) → full weight
        ma = next(r for r in result if r.signal_type == "ma_stack")
        rsi = next(r for r in result if r.signal_type == "rsi")

        assert ma.penalized_log_lr == 0.8  # rank 0: 0.5^0 = 1.0
        assert abs(rsi.penalized_log_lr - 0.5 * 0.5) < 1e-10  # rank 1: 0.5^1 = 0.5

    def test_three_structure_signals_decay(self) -> None:
        """Three signals in same cluster: 1.0, 0.5, 0.25 decay."""
        lrs = [
            _lr("fibonacci", "structure", log_lr=1.0),
            _lr("elliott_wave", "structure", log_lr=0.7),
            _lr("cup_handle", "structure", log_lr=0.3),
        ]
        result = apply_correlation_penalty(lrs)

        fib = next(r for r in result if r.signal_type == "fibonacci")
        ew = next(r for r in result if r.signal_type == "elliott_wave")
        ch = next(r for r in result if r.signal_type == "cup_handle")

        assert fib.penalized_log_lr == 1.0  # rank 0: 1.0
        assert abs(ew.penalized_log_lr - 0.7 * 0.5) < 1e-10  # rank 1: 0.5
        assert abs(ch.penalized_log_lr - 0.3 * 0.25) < 1e-10  # rank 2: 0.25

    def test_ranking_by_absolute_log_lr(self) -> None:
        """Ranking uses abs(log_lr), so a negative LR with large magnitude ranks first."""
        lrs = [
            _lr("ma_stack", "momentum", log_lr=0.3),
            _lr("rsi", "momentum", log_lr=-0.9),  # abs = 0.9, strongest
        ]
        result = apply_correlation_penalty(lrs)

        rsi = next(r for r in result if r.signal_type == "rsi")
        ma = next(r for r in result if r.signal_type == "ma_stack")

        # RSI is strongest by abs → full weight
        assert rsi.penalized_log_lr == -0.9
        # MA is second → 0.5 decay
        assert abs(ma.penalized_log_lr - 0.3 * 0.5) < 1e-10

    def test_decay_reduces_penalized_log_lr_magnitude(self) -> None:
        """Penalized log_lr magnitude is always <= original for non-strongest."""
        lrs = [
            _lr("ma_stack", "momentum", log_lr=0.8),
            _lr("rsi", "momentum", log_lr=0.6),
        ]
        result = apply_correlation_penalty(lrs)

        rsi = next(r for r in result if r.signal_type == "rsi")
        assert abs(rsi.penalized_log_lr) < abs(rsi.log_lr)


# ===========================================================================
# 3. apply_correlation_penalty — cross-cluster independence (Requirement 7.3)
# ===========================================================================


class TestCrossClusterIndependence:
    """Signals from different clusters receive no penalty."""

    def test_different_clusters_no_penalty(self) -> None:
        """Each signal in its own cluster → all at full weight."""
        lrs = [
            _lr("ma_stack", "momentum", log_lr=0.8),
            _lr("fibonacci", "structure", log_lr=0.7),
            _lr("atr", "volatility", log_lr=0.5),
            _lr("valuation", "fundamentals", log_lr=0.3),
        ]
        result = apply_correlation_penalty(lrs)

        for r in result:
            assert r.penalized_log_lr == r.log_lr, (
                f"{r.signal_type}: penalized_log_lr should equal log_lr "
                f"when alone in cluster"
            )

    def test_mixed_clusters_only_same_cluster_penalized(self) -> None:
        """Two momentum + one structure: only momentum signals get decay."""
        lrs = [
            _lr("ma_stack", "momentum", log_lr=0.8),
            _lr("rsi", "momentum", log_lr=0.5),
            _lr("fibonacci", "structure", log_lr=0.6),
        ]
        result = apply_correlation_penalty(lrs)

        ma = next(r for r in result if r.signal_type == "ma_stack")
        rsi = next(r for r in result if r.signal_type == "rsi")
        fib = next(r for r in result if r.signal_type == "fibonacci")

        # Momentum cluster: ma_stack full, rsi decayed
        assert ma.penalized_log_lr == 0.8
        assert abs(rsi.penalized_log_lr - 0.5 * 0.5) < 1e-10

        # Structure cluster: fibonacci alone → no penalty
        assert fib.penalized_log_lr == 0.6


# ===========================================================================
# 4. apply_correlation_penalty — single-signal clusters (Requirement 7.4)
# ===========================================================================


class TestSingleSignalCluster:
    """Single-signal clusters receive no penalty."""

    def test_single_signal_no_penalty(self) -> None:
        """One signal in a cluster → penalized_log_lr == log_lr."""
        lrs = [_lr("fibonacci", "structure", log_lr=0.9)]
        result = apply_correlation_penalty(lrs)

        assert len(result) == 1
        assert result[0].penalized_log_lr == 0.9

    def test_multiple_single_signal_clusters(self) -> None:
        """Multiple clusters each with one signal → no penalties anywhere."""
        lrs = [
            _lr("rsi", "momentum", log_lr=0.4),
            _lr("fibonacci", "structure", log_lr=0.6),
        ]
        result = apply_correlation_penalty(lrs)

        for r in result:
            assert r.penalized_log_lr == r.log_lr


# ===========================================================================
# 5. Edge cases
# ===========================================================================


class TestEdgeCases:
    """Edge cases: empty input, zero log_lr, original order preserved."""

    def test_empty_input_returns_empty(self) -> None:
        """Empty list → empty list."""
        assert apply_correlation_penalty([]) == []

    def test_zero_log_lr_no_effect(self) -> None:
        """log_lr = 0 → penalized_log_lr = 0 regardless of rank."""
        lrs = [
            _lr("ma_stack", "momentum", log_lr=0.5),
            _lr("rsi", "momentum", log_lr=0.0),
        ]
        result = apply_correlation_penalty(lrs)

        rsi = next(r for r in result if r.signal_type == "rsi")
        assert rsi.penalized_log_lr == 0.0

    def test_original_order_preserved(self) -> None:
        """Output list preserves the original input order."""
        lrs = [
            _lr("rsi", "momentum", log_lr=0.3),
            _lr("fibonacci", "structure", log_lr=0.9),
            _lr("ma_stack", "momentum", log_lr=0.8),
        ]
        result = apply_correlation_penalty(lrs)

        assert result[0].signal_type == "rsi"
        assert result[1].signal_type == "fibonacci"
        assert result[2].signal_type == "ma_stack"

    def test_original_objects_not_mutated(self) -> None:
        """Input LikelihoodRatio objects are not modified in place."""
        original = _lr("ma_stack", "momentum", log_lr=0.8)
        lrs = [
            original,
            _lr("rsi", "momentum", log_lr=0.5),
        ]
        apply_correlation_penalty(lrs)

        # Original object should still have its initial penalized_log_lr
        assert original.penalized_log_lr == 0.8

    def test_negative_log_lr_decay(self) -> None:
        """Negative log_lr values are decayed correctly (toward zero)."""
        lrs = [
            _lr("ma_stack", "momentum", log_lr=-0.8),
            _lr("rsi", "momentum", log_lr=-0.4),
        ]
        result = apply_correlation_penalty(lrs)

        ma = next(r for r in result if r.signal_type == "ma_stack")
        rsi = next(r for r in result if r.signal_type == "rsi")

        # ma_stack strongest by abs → full weight
        assert ma.penalized_log_lr == -0.8
        # rsi second → 0.5 decay
        assert abs(rsi.penalized_log_lr - (-0.4 * 0.5)) < 1e-10

    def test_lr_field_unchanged_by_penalty(self) -> None:
        """The raw lr field is preserved unchanged through penalty."""
        lr_val = math.exp(0.5)
        lrs = [
            _lr("ma_stack", "momentum", log_lr=0.8),
            LikelihoodRatio(
                signal_type="rsi",
                cluster="momentum",
                lr=lr_val,
                log_lr=0.5,
                penalized_log_lr=0.5,
                hit_rate=0.6,
                strength=0.7,
            ),
        ]
        result = apply_correlation_penalty(lrs)

        rsi = next(r for r in result if r.signal_type == "rsi")
        assert rsi.lr == lr_val
        assert rsi.hit_rate == 0.6
        assert rsi.strength == 0.7
