#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""Reward momentum -- Item 136: consecutive success biology fix (P33).

Extracted from reward_signal.py to keep that file under 300 lines.
Biologically: consecutive expected successes DECREASE RPE (prediction error → 0).
A surprise failure after a success streak gets a large NEGATIVE RPE amplification.
"""
import logging

log = logging.getLogger("pulse.reward_momentum")

# ---------------------------------------------------------------------------
# Item 136: Reward momentum -- biologically correct (P33)
# ---------------------------------------------------------------------------
_CONSECUTIVE_SUCCESSES: dict = {}  # {domain: count}
_MOMENTUM_RPE_STEP = 0.05         # Step size per consecutive success (shrinks toward 0)
_MOMENTUM_MAX_BONUS = 0.30        # Maximum RPE reduction from baseline per streak
_FAILURE_AMPLIFY_FACTOR = 2.0     # Surprise failure after streak gets amplified RPE


def record_success_momentum(domain: str, success: bool) -> float:
    """Track consecutive successes per domain and return RPE momentum delta.

    Biology: after N consecutive expected successes the RPE signal shrinks toward 0
    (prediction error goes to 0 because the outcome is fully anticipated).
    A surprise failure after a streak gets a NEGATIVE amplification (large prediction error).

    Returns the momentum delta: negative (reduces RPE) on success streaks,
    negative-amplified on failure-after-streak, 0.0 on isolated failure.
    """
    try:
        key = domain or "general"
        if success:
            streak_before = _CONSECUTIVE_SUCCESSES.get(key, 0)
            _CONSECUTIVE_SUCCESSES[key] = streak_before + 1
            streak = _CONSECUTIVE_SUCCESSES[key]
            # Each consecutive success shrinks RPE toward 0 (subtract from RPE)
            reduction = min(_MOMENTUM_MAX_BONUS, (streak - 1) * _MOMENTUM_RPE_STEP)
            return round(-reduction, 4)
        else:
            streak = _CONSECUTIVE_SUCCESSES.get(key, 0)
            _CONSECUTIVE_SUCCESSES[key] = 0
            if streak >= 2:
                # Surprise failure after a run — amplify negative RPE
                amplification = round(-_FAILURE_AMPLIFY_FACTOR * _MOMENTUM_RPE_STEP * streak, 4)
                return max(-_MOMENTUM_MAX_BONUS * _FAILURE_AMPLIFY_FACTOR, amplification)
            return 0.0
    except Exception as e:
        log.debug("record_success_momentum failed: %s", e)
        return 0.0


def get_domain_streak(domain: str) -> int:
    """Return current consecutive success streak for a domain."""
    try:
        return _CONSECUTIVE_SUCCESSES.get(domain or "general", 0)
    except Exception:
        return 0


def reset_domain_streak(domain: str) -> None:
    """Manually reset streak (e.g. after agent restart)."""
    try:
        _CONSECUTIVE_SUCCESSES[domain or "general"] = 0
    except Exception:
        pass
