#!/usr/bin/env python3
"""
Audit: Reward Signal RPE Computation for Surprising Outcomes.

Covers gaps identified in the reward signal pipeline audit (t_1773842941_010):
1. compute_rpe: clamping at extreme expected rates (0.0, 1.0)
2. compute_rpe: maximum-surprise magnitude boundaries (±0.95 after clamping)
3. assign_credit: loss aversion 2x amplification for negative RPE (surprising failure)
4. assign_credit: edge cases — zero rpe_delta, empty item list
5. assign_credit: MIN_CREDIT threshold excludes low-credit tail items
6. Fallback path discrepancy: _record_reward uses raw confidence, not RPE
7. Volatility: RPE variance signal computation path
8. record_task_outcome: full reward chain fires and returns valid RewardSignal

Brain anti-pattern context:
- Reward signals split between reward_signal.py and prediction_feedback.py
- Loss aversion doubling logic lacked test coverage
- RPE credit magnitude logic unverified for surprising (extreme) outcomes
"""
import json
import math
import tempfile
from pathlib import Path
from unittest.mock import patch

import pytest

from pulse.core.reward.reward_signal import (
    CREDIT_DECAY,
    CONTRADICT_DELTA,
    LOSS_AVERSION_RATIO,
    MIN_CREDIT,
    REINFORCE_DELTA,
    RewardSignal,
    assign_credit,
    compute_rpe,
)


# ── 1. compute_rpe: clamping and surprise magnitude ───────────────────────────

class TestComputeRPEClamping:
    """Expected success rate is clamped to [0.05, 0.95] before subtraction."""

    def test_zero_expected_clamped_to_005_on_success(self):
        """Expected=0.0 clamped to 0.05 → RPE = 1.0 - 0.05 = 0.95 (max positive surprise)."""
        rpe = compute_rpe(success=True, expected_success_rate=0.0)
        assert rpe == pytest.approx(0.95, abs=1e-9), (
            f"Expected 0.95 (max positive surprise), got {rpe}"
        )

    def test_unity_expected_clamped_to_095_on_failure(self):
        """Expected=1.0 clamped to 0.95 → RPE = 0.0 - 0.95 = -0.95 (max negative surprise)."""
        rpe = compute_rpe(success=False, expected_success_rate=1.0)
        assert rpe == pytest.approx(-0.95, abs=1e-9), (
            f"Expected -0.95 (max negative surprise), got {rpe}"
        )

    def test_below_zero_expected_clamped(self):
        """Negative expected_success_rate must be clamped to 0.05."""
        rpe = compute_rpe(success=True, expected_success_rate=-0.5)
        # Same as expected=0.05
        assert rpe == pytest.approx(0.95, abs=1e-9)

    def test_above_unity_expected_clamped(self):
        """expected_success_rate > 1.0 must be clamped to 0.95."""
        rpe = compute_rpe(success=False, expected_success_rate=1.5)
        assert rpe == pytest.approx(-0.95, abs=1e-9)

    def test_rpe_range_bounded_by_negative_095_to_095(self):
        """RPE must always lie within [-0.95, 0.95] — the docstring says ~[-1, 1]
        but clamping means the actual range is tighter."""
        for success in (True, False):
            for rate in (0.0, 0.05, 0.1, 0.5, 0.9, 0.95, 1.0):
                rpe = compute_rpe(success=success, expected_success_rate=rate)
                assert -0.95 <= rpe <= 0.95, (
                    f"RPE out of [-0.95, 0.95]: success={success}, rate={rate}, rpe={rpe}"
                )


class TestComputeRPESurpriseDirection:
    """Surprising outcomes produce RPE in the correct direction."""

    def test_surprising_success_positive_rpe(self):
        """Low expected rate + success = positive RPE (unexpectedly good)."""
        rpe = compute_rpe(success=True, expected_success_rate=0.1)
        assert rpe > 0, f"Surprising success should produce positive RPE, got {rpe}"
        assert rpe == pytest.approx(0.9, abs=1e-9)

    def test_surprising_failure_negative_rpe(self):
        """High expected rate + failure = negative RPE (unexpectedly bad)."""
        rpe = compute_rpe(success=False, expected_success_rate=0.9)
        assert rpe < 0, f"Surprising failure should produce negative RPE, got {rpe}"
        assert rpe == pytest.approx(-0.9, abs=1e-9)

    def test_expected_success_small_positive_rpe(self):
        """Expected=0.95 + success → small positive RPE (not surprising)."""
        rpe = compute_rpe(success=True, expected_success_rate=0.95)
        assert 0 < rpe < 0.1, f"Expected small positive RPE for expected success, got {rpe}"

    def test_expected_failure_small_negative_rpe(self):
        """Expected=0.05 + failure → small negative RPE (not surprising)."""
        rpe = compute_rpe(success=False, expected_success_rate=0.05)
        assert -0.1 < rpe < 0, f"Expected small negative RPE for expected failure, got {rpe}"

    def test_fifty_fifty_success_half_positive(self):
        """50/50 expected + success → RPE = 0.5."""
        rpe = compute_rpe(success=True, expected_success_rate=0.5)
        assert rpe == pytest.approx(0.5, abs=1e-9)

    def test_fifty_fifty_failure_half_negative(self):
        """50/50 expected + failure → RPE = -0.5."""
        rpe = compute_rpe(success=False, expected_success_rate=0.5)
        assert rpe == pytest.approx(-0.5, abs=1e-9)

    def test_surprise_magnitude_proportional_to_gap(self):
        """Larger gap between actual and expected → larger |RPE|."""
        rpe_low_surprise = abs(compute_rpe(success=True, expected_success_rate=0.8))
        rpe_high_surprise = abs(compute_rpe(success=True, expected_success_rate=0.2))
        assert rpe_high_surprise > rpe_low_surprise, (
            "Larger prediction gap must produce larger |RPE|"
        )


# ── 2. assign_credit: loss aversion amplification ─────────────────────────────

class TestAssignCreditLossAversion:
    """Negative RPE (surprising failure) applies LOSS_AVERSION_RATIO × amplification."""

    def test_negative_rpe_2x_positive_rpe_magnitude(self):
        """assign_credit(-0.5) must produce exactly 2x credit magnitude vs assign_credit(+0.5).

        Kahneman & Tversky: losses are amplified by LOSS_AVERSION_RATIO (2.0).
        """
        positive_credits = assign_credit(["x"], rpe_delta=0.5)
        negative_credits = assign_credit(["x"], rpe_delta=-0.5)
        pos_mag = abs(positive_credits["x"])
        neg_mag = abs(negative_credits["x"])
        assert neg_mag / pos_mag == pytest.approx(LOSS_AVERSION_RATIO, abs=1e-6), (
            f"Loss aversion ratio must be {LOSS_AVERSION_RATIO}x, "
            f"got {neg_mag / pos_mag:.4f}"
        )

    def test_negative_rpe_credit_is_negative(self):
        """Negative RPE (surprising failure) must produce negative credits (blame)."""
        credits = assign_credit(["a", "b"], rpe_delta=-0.6)
        for item_id, credit in credits.items():
            assert credit < 0, f"Negative RPE must produce negative credit for {item_id}: {credit}"

    def test_positive_rpe_credit_is_positive(self):
        """Positive RPE (surprising success) must produce positive credits (reinforce)."""
        credits = assign_credit(["a", "b"], rpe_delta=0.6)
        for item_id, credit in credits.items():
            assert credit > 0, f"Positive RPE must produce positive credit for {item_id}: {credit}"

    def test_loss_aversion_amplification_for_max_surprise(self):
        """Max surprise failure: RPE=-0.95 → magnitude 0.95 × 2.0 = 1.9 for first item."""
        credits = assign_credit(["only"], rpe_delta=-0.95)
        expected_magnitude = 0.95 * LOSS_AVERSION_RATIO  # 1.9
        assert abs(credits["only"]) == pytest.approx(expected_magnitude, abs=1e-4), (
            f"Max-surprise failure credit must be {expected_magnitude}, got {abs(credits['only'])}"
        )

    def test_surprising_failure_larger_than_surprising_success_at_same_rpe(self):
        """For same |rpe|, failure credit magnitude must exceed success credit magnitude."""
        rpe_mag = 0.7
        pos = assign_credit(["item"], rpe_delta=rpe_mag)
        neg = assign_credit(["item"], rpe_delta=-rpe_mag)
        assert abs(neg["item"]) > abs(pos["item"]), (
            "Surprising failure must receive stronger credit signal than surprising success"
        )


# ── 3. assign_credit: edge cases ──────────────────────────────────────────────

class TestAssignCreditEdgeCases:
    """Edge cases: zero rpe_delta, empty items, MIN_CREDIT threshold."""

    def test_zero_rpe_returns_empty_dict(self):
        """rpe_delta=0 means no surprise → no credit assigned."""
        credits = assign_credit(["a", "b", "c"], rpe_delta=0.0)
        assert credits == {}, f"Zero RPE should return empty dict, got {credits}"

    def test_empty_items_returns_empty_dict(self):
        """No retrieved items → no credit to assign."""
        credits = assign_credit([], rpe_delta=0.8)
        assert credits == {}, f"Empty items should return empty dict, got {credits}"

    def test_positional_decay_monotonically_decreasing(self):
        """Earlier items in retrieval list must receive more credit (higher |credit|)."""
        items = ["a", "b", "c", "d", "e"]
        credits = assign_credit(items, rpe_delta=0.9)
        magnitudes = [abs(credits[i]) for i in items if i in credits]
        for j in range(len(magnitudes) - 1):
            assert magnitudes[j] >= magnitudes[j + 1], (
                f"Credit must decrease with position: index {j}={magnitudes[j]} "
                f"vs index {j+1}={magnitudes[j+1]}"
            )

    def test_positional_decay_uses_credit_decay_constant(self):
        """Second item credit / first item credit should equal CREDIT_DECAY (0.95)."""
        credits = assign_credit(["first", "second"], rpe_delta=0.6)
        ratio = abs(credits["second"]) / abs(credits["first"])
        assert ratio == pytest.approx(CREDIT_DECAY, abs=1e-6), (
            f"Positional decay ratio must be {CREDIT_DECAY}, got {ratio:.6f}"
        )

    def test_min_credit_threshold_excludes_tail_items(self):
        """Items whose computed credit falls below MIN_CREDIT must be excluded.

        With rpe_delta=0.06 (positive, no amplification), CREDIT_DECAY^4 * 0.06 ≈ 0.0488 < MIN_CREDIT=0.05.
        So item at index 4 (and beyond) should be excluded.
        """
        items = [f"item{i}" for i in range(10)]
        credits = assign_credit(items, rpe_delta=0.06)
        # All included items must have |credit| >= MIN_CREDIT
        for item_id, credit in credits.items():
            assert abs(credit) >= MIN_CREDIT, (
                f"{item_id} credit {credit} is below MIN_CREDIT={MIN_CREDIT}"
            )
        # There must be fewer credits than items (threshold excludes tail)
        assert len(credits) < len(items), (
            f"MIN_CREDIT threshold should exclude some tail items, but all {len(items)} included"
        )

    def test_single_item_gets_full_magnitude(self):
        """Single retrieved item gets full (undecayed) credit."""
        credits = assign_credit(["solo"], rpe_delta=0.7)
        # Positional decay for index 0 = CREDIT_DECAY^0 = 1.0
        assert abs(credits["solo"]) == pytest.approx(0.7, abs=1e-4), (
            f"Solo item should get undecayed credit=0.7, got {credits['solo']}"
        )


# ── 4. Fallback path discrepancy audit ────────────────────────────────────────

class TestFallbackRewardDiscrepancy:
    """Audit: _record_reward (fallback) uses raw confidence, not RPE.

    When reward_signal.py fails to import, prediction_feedback._record_reward()
    stores prediction_confidence as reward_score. This is NOT an RPE:
    - It ignores loss aversion
    - It writes to state/reward_signals.json (not state/reward/reward_signals.json)
    - It doesn't go through the unified reward pipeline
    This is a documented anti-pattern — test verifies the discrepancy exists.
    """

    def test_fallback_uses_raw_confidence_not_rpe(self):
        """_record_reward uses prediction_confidence as reward_score, not RPE.

        Audit finding: the fallback reward score is the raw prediction confidence
        (e.g., 0.8), NOT compute_rpe(success=True, expected=0.8) which would be
        compute_rpe(True, 0.8) = 0.2. A confidence of 0.8 ≠ RPE of 0.2.
        """
        import inspect
        import pulse.daemons.neural.prediction_feedback as fb

        source = inspect.getsource(fb._record_reward)
        # Verify fallback stores raw confidence, not RPE
        assert "reward_score = round(prediction_confidence" in source, (
            "_record_reward must store raw prediction_confidence as reward_score "
            "(this is the known discrepancy — raw confidence ≠ RPE)"
        )
        # Verify no RPE computation in fallback
        assert "compute_rpe" not in source, (
            "_record_reward fallback does NOT call compute_rpe (documented anti-pattern)"
        )

    def test_fallback_writes_to_different_file_than_main(self):
        """Fallback path writes to state/reward_signals.json; main path writes to state/reward/reward_signals.json."""
        import pulse.daemons.neural.prediction_feedback as fb
        from pulse.core.reward.reward_signal import REWARD_SIGNALS_FILE as main_file

        fallback_file = fb.REWARD_SIGNALS_FILE
        # They should differ: fallback is state/reward_signals.json,
        # main is state/reward/reward_signals.json
        assert fallback_file != main_file, (
            f"Fallback writes to '{fallback_file}', main writes to '{main_file}'. "
            "They should be different — this is the documented split anti-pattern."
        )
        # Fallback is in state root, main is in state/reward/ subdirectory
        assert fallback_file.parent.name != "reward", (
            "Fallback file must be in state/ root, not state/reward/ (confirms the split)"
        )
        assert main_file.parent.name == "reward", (
            "Main reward file must be in state/reward/ subdirectory"
        )

    def test_main_path_preferred_over_fallback(self):
        """record_prediction_outcome calls record_task_outcome (main path), not _record_reward directly."""
        import inspect
        import pulse.daemons.neural.prediction_feedback as fb

        source = inspect.getsource(fb.record_prediction_outcome)
        # Main path is attempted first
        assert "record_task_outcome" in source, (
            "record_prediction_outcome must try record_task_outcome (unified pipeline) first"
        )
        # Fallback is only on ImportError
        assert "ImportError" in source, (
            "_record_reward fallback must only trigger on ImportError"
        )
        assert "except ImportError" in source, (
            "Fallback must be guarded by ImportError, not a broad except"
        )


# ── 5. Volatility RPE variance signal ─────────────────────────────────────────

class TestVolatilityRPEVariance:
    """compute_volatility() uses RPE variance from recent reward signals."""

    def test_high_rpe_variance_increases_multiplier(self):
        """High RPE variance (mixed positive and negative) → higher volatility multiplier.

        volatility.py does `from pulse.core.reward.reward_signal import REWARD_SIGNALS_FILE`
        inside the function body, so we patch the source module attribute.
        """
        from pulse.core.reward.volatility import compute_volatility

        high_var_signals = [
            {"rpe_delta": 0.9 if i % 2 == 0 else -0.9, "success": i % 2 == 0}
            for i in range(20)
        ]
        low_var_signals = [
            {"rpe_delta": 0.1, "success": True}
            for _ in range(20)
        ]

        with tempfile.TemporaryDirectory() as tmp:
            high_file = Path(tmp) / "high_var.json"
            low_file = Path(tmp) / "low_var.json"
            vol_file = Path(tmp) / "vol.json"
            high_file.write_text(json.dumps(high_var_signals))
            low_file.write_text(json.dumps(low_var_signals))

            # volatility.py imports REWARD_SIGNALS_FILE from reward_signal inside the fn
            # → must patch the source module (pulse.core.reward_signal)
            with patch("pulse.core.reward.reward_signal.REWARD_SIGNALS_FILE", high_file), \
                 patch("pulse.core.reward.volatility.VOLATILITY_FILE", vol_file), \
                 patch("pulse.core.brain_utils.read_jsonl_entries", return_value=[]), \
                 patch("pulse.core.brain_db.get_conn", side_effect=Exception("no db")):
                high_mult = compute_volatility()

            with patch("pulse.core.reward.reward_signal.REWARD_SIGNALS_FILE", low_file), \
                 patch("pulse.core.reward.volatility.VOLATILITY_FILE", vol_file), \
                 patch("pulse.core.brain_utils.read_jsonl_entries", return_value=[]), \
                 patch("pulse.core.brain_db.get_conn", side_effect=Exception("no db")):
                low_mult = compute_volatility()

        assert high_mult > low_mult, (
            f"High RPE variance should produce higher multiplier ({high_mult}) "
            f"than low variance ({low_mult})"
        )

    def test_rpe_variance_formula(self):
        """RPE signal = min(1.5, 0.5 + var_rpe * 5) — verify math directly."""
        # Perfect certainty (all same RPE): variance=0 → rpe_signal=0.5
        rpes_constant = [0.5] * 20
        mean = sum(rpes_constant) / len(rpes_constant)
        var = sum((r - mean) ** 2 for r in rpes_constant) / len(rpes_constant)
        rpe_signal = min(1.5, 0.5 + var * 5)
        assert rpe_signal == pytest.approx(0.5, abs=1e-9), (
            f"Zero variance should produce signal=0.5, got {rpe_signal}"
        )

        # Maximum variance (alternating ±1): variance ≈ 1.0 → rpe_signal = min(1.5, 5.5) = 1.5
        rpes_max_var = [1.0 if i % 2 == 0 else -1.0 for i in range(20)]
        mean2 = sum(rpes_max_var) / len(rpes_max_var)
        var2 = sum((r - mean2) ** 2 for r in rpes_max_var) / len(rpes_max_var)
        rpe_signal2 = min(1.5, 0.5 + var2 * 5)
        assert rpe_signal2 == pytest.approx(1.5, abs=1e-6), (
            f"Max variance should produce signal=1.5 (capped), got {rpe_signal2}"
        )

    def test_volatility_clamped_to_multiplier_range(self):
        """Volatility multiplier is always in [MIN_MULTIPLIER, MAX_MULTIPLIER]."""
        from pulse.core.reward.volatility import MIN_MULTIPLIER, MAX_MULTIPLIER, compute_volatility

        with tempfile.TemporaryDirectory() as tmp:
            extreme_signals = [
                {"rpe_delta": 0.95 if i % 2 == 0 else -0.95}
                for i in range(20)
            ]
            sig_file = Path(tmp) / "extreme.json"
            sig_file.write_text(json.dumps(extreme_signals))

            # Patch source module — volatility imports REWARD_SIGNALS_FILE from there
            with patch("pulse.core.reward.reward_signal.REWARD_SIGNALS_FILE", sig_file), \
                 patch("pulse.core.reward.volatility.VOLATILITY_FILE", Path(tmp) / "v.json"), \
                 patch("pulse.core.brain_utils.read_jsonl_entries", return_value=[]), \
                 patch("pulse.core.brain_db.get_conn", side_effect=Exception("no db")):
                mult = compute_volatility()

        assert MIN_MULTIPLIER <= mult <= MAX_MULTIPLIER, (
            f"Volatility multiplier {mult} out of range [{MIN_MULTIPLIER}, {MAX_MULTIPLIER}]"
        )


# ── 6. Full reward chain fires ────────────────────────────────────────────────

class TestRewardChainFires:
    """record_task_outcome must execute all pipeline steps and return valid RewardSignal."""

    def test_reward_chain_returns_reward_signal(self):
        """record_task_outcome returns a RewardSignal dataclass with correct fields."""
        from pulse.core.reward.reward_signal import record_task_outcome

        with tempfile.TemporaryDirectory() as tmp:
            tmp_reward_dir = Path(tmp) / "reward"
            with patch("pulse.core.reward.reward_signal._REWARD_DIR", tmp_reward_dir), \
                 patch("pulse.core.reward.reward_signal.REWARD_SIGNALS_FILE",
                        tmp_reward_dir / "reward_signals.json"), \
                 patch("pulse.core.reward.reward_signal.CREDIT_LEDGER_FILE",
                        tmp_reward_dir / "credit_ledger.json"), \
                 patch("pulse.core.reward.reward_signal._get_expected_success_rate",
                        return_value=0.3), \
                 patch("pulse.core.reward.reward_signal.propagate_reward", return_value=2):
                signal = record_task_outcome(
                    task_id="test_task_001", agent="test_agent",
                    success=True, domain="testing",
                    retrieved_item_ids=["L-001", "L-002", "L-003"],
                )

        assert isinstance(signal, RewardSignal), f"Must return RewardSignal, got {type(signal)}"
        assert signal.task_id == "test_task_001"
        assert signal.agent == "test_agent"
        assert signal.success is True
        assert signal.domain == "testing"
        # Surprising success (expected=0.3): RPE = 1.0 - 0.3 = 0.7
        assert signal.rpe_delta == pytest.approx(0.7, abs=1e-4), (
            f"RPE should be 0.7 (surprising success with expected=0.3), got {signal.rpe_delta}"
        )
        assert len(signal.items_credited) > 0, "Credits must be assigned for retrieved items"

    def test_reward_signal_logged_to_file(self):
        """record_task_outcome must persist the signal to REWARD_SIGNALS_FILE."""
        from pulse.core.reward.reward_signal import record_task_outcome

        with tempfile.TemporaryDirectory() as tmp:
            tmp_reward_dir = Path(tmp) / "reward"
            signals_file = tmp_reward_dir / "reward_signals.json"
            with patch("pulse.core.reward.reward_signal._REWARD_DIR", tmp_reward_dir), \
                 patch("pulse.core.reward.reward_signal.REWARD_SIGNALS_FILE", signals_file), \
                 patch("pulse.core.reward.reward_signal.CREDIT_LEDGER_FILE",
                        tmp_reward_dir / "credit_ledger.json"), \
                 patch("pulse.core.reward.reward_signal._get_expected_success_rate",
                        return_value=0.5), \
                 patch("pulse.core.reward.reward_signal.propagate_reward", return_value=0):
                record_task_outcome(
                    task_id="logged_task", agent="test_agent",
                    success=False, domain="testing",
                )

            # Assertions must be INSIDE the tempdir context (dir is deleted on exit)
            assert signals_file.exists(), "REWARD_SIGNALS_FILE must be written after record_task_outcome"
            data = json.loads(signals_file.read_text())
            assert isinstance(data, list) and len(data) == 1, "Signal must be appended to list"
            assert data[0]["task_id"] == "logged_task"
            assert data[0]["success"] is False
            # Surprising failure (expected=0.5): RPE = 0.0 - 0.5 = -0.5
            assert data[0]["rpe_delta"] == pytest.approx(-0.5, abs=1e-4)

    def test_surprising_failure_credits_have_loss_aversion(self):
        """End-to-end: surprising failure assigns 2x larger credits than surprising success."""
        rpe_fail = compute_rpe(success=False, expected_success_rate=0.8)  # = -0.8
        rpe_succ = compute_rpe(success=True, expected_success_rate=0.2)   # = +0.8

        items = ["item_A"]
        credits_fail = assign_credit(items, rpe_fail)
        credits_succ = assign_credit(items, rpe_succ)

        assert abs(rpe_fail) == abs(rpe_succ), "Both have same |RPE| = 0.8"
        fail_mag = abs(credits_fail["item_A"])
        succ_mag = abs(credits_succ["item_A"])
        assert fail_mag / succ_mag == pytest.approx(LOSS_AVERSION_RATIO, abs=1e-6), (
            f"Surprising failure credit ({fail_mag}) must be {LOSS_AVERSION_RATIO}x "
            f"surprising success credit ({succ_mag})"
        )
