# Feature: dual-pipeline-signal-engine, Property: Correlation penalty reduces confidence
"""Property-based tests for the signal correlation penalty.

Feature: dual-pipeline-signal-engine

Tests that the within-cluster correlation penalty always reduces (or maintains)
the posterior probability compared to the unpenalized posterior.  Correlated
signals within the same cluster receive exponential decay (0.5^(n-1)), so the
penalized Σ log(LR_i) is always <= the unpenalized Σ log(LR_i) in absolute
magnitude, which means the posterior moves less from the prior.

Requirements: 7.5, 17.4
"""
from __future__ import annotations

import math

from hypothesis import given, settings
from hypothesis import strategies as st

from services.signal_engine.correlation import apply_correlation_penalty
from services.signal_engine.models import LikelihoodRatio
from services.signal_engine.probabilistic import _logit, _sigmoid

# ---------------------------------------------------------------------------
# Hypothesis strategies
# ---------------------------------------------------------------------------

# Clusters that have multiple signal types mapped to them
_CLUSTER_SIGNAL_TYPES: dict[str, list[str]] = {
    "momentum": ["ma_stack", "rsi"],
    "structure": ["fibonacci", "elliott_wave", "cup_handle"],
    "volatility": ["atr", "bollinger"],
    "fundamentals": ["valuation", "earnings", "macro"],
}

# Pick a cluster that has at least 2 signal types
_cluster_with_types = st.sampled_from(
    [(cluster, types) for cluster, types in _CLUSTER_SIGNAL_TYPES.items() if len(types) >= 2]
)

# Positive log-LR values (bullish signals) — ensures same direction within cluster
_positive_log_lr = st.floats(
    min_value=0.01, max_value=5.0, allow_nan=False, allow_infinity=False,
)

# Prior probability
_prior_prob = st.floats(
    min_value=0.01, max_value=0.99, allow_nan=False, allow_infinity=False,
)


@st.composite
def _correlated_lr_set(draw: st.DrawFn) -> list[LikelihoodRatio]:
    """Generate a list of LikelihoodRatio objects with at least 2 in the same cluster.

    All signals within the chosen cluster have positive log_lr (bullish direction)
    so the penalty effect is clearly measurable.
    """
    cluster, signal_types = draw(_cluster_with_types)

    # Draw at least 2 signals from the same cluster
    n_correlated = draw(st.integers(min_value=2, max_value=len(signal_types)))
    chosen_types = draw(
        st.lists(
            st.sampled_from(signal_types),
            min_size=n_correlated,
            max_size=n_correlated,
            unique=True,
        )
    )

    lrs: list[LikelihoodRatio] = []
    for sig_type in chosen_types:
        log_lr = draw(_positive_log_lr)
        lr_val = math.exp(log_lr)
        lrs.append(
            LikelihoodRatio(
                signal_type=sig_type,
                cluster=cluster,
                lr=lr_val,
                log_lr=log_lr,
                penalized_log_lr=log_lr,  # unpenalized initially
                hit_rate=0.6,
                strength=0.5,
            )
        )

    return lrs


# ---------------------------------------------------------------------------
# Property: Correlation penalty reduces confidence
# Validates: Requirements 7.5, 17.4
# ---------------------------------------------------------------------------


@given(lrs=_correlated_lr_set(), p_prior=_prior_prob)
@settings(max_examples=100)
def test_penalized_posterior_leq_unpenalized(
    lrs: list[LikelihoodRatio], p_prior: float,
) -> None:
    """**Validates: Requirements 7.5, 17.4**

    For any signal set with at least 2 correlated signals in the same
    cluster, the penalized posterior SHALL be <= the unpenalized posterior.

    The penalty reduces the magnitude of Σ penalized_log_lr relative to
    Σ log_lr, so the posterior moves less from the prior.
    """
    # Unpenalized posterior: use raw log_lr values
    logit_prior = _logit(p_prior)
    sum_unpenalized = sum(lr.log_lr for lr in lrs)
    p_unpenalized = _sigmoid(logit_prior + sum_unpenalized)

    # Apply correlation penalty
    penalized_lrs = apply_correlation_penalty(lrs)

    # Penalized posterior: use penalized_log_lr values
    sum_penalized = sum(lr.penalized_log_lr for lr in penalized_lrs)
    p_penalized = _sigmoid(logit_prior + sum_penalized)

    # Since all log_lr values are positive (bullish), the penalty reduces
    # the sum, which means the penalized posterior is <= unpenalized posterior
    assert p_penalized <= p_unpenalized + 1e-12, (
        f"Penalized posterior {p_penalized} > unpenalized {p_unpenalized}. "
        f"Prior={p_prior}, Σ_raw={sum_unpenalized}, Σ_penalized={sum_penalized}"
    )

    # Also verify the penalized sum of log-LRs is <= the unpenalized sum
    assert sum_penalized <= sum_unpenalized + 1e-12, (
        f"Penalized Σ log_lr ({sum_penalized}) > unpenalized ({sum_unpenalized})"
    )
