#!/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
from pathlib import Path

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

# ---------------------------------------------------------------------------
# Item 136: Reward momentum -- biologically correct (P33)
# ---------------------------------------------------------------------------
_CONSECUTIVE_SUCCESSES: dict = {}  # {domain: count}
_STATE_LOADED = False

def _state_file() -> Path:
    from pulse.core.paths import get_state_dir
    return get_state_dir() / "reward" / "reward_momentum_state.json"

def _load_state() -> None:
    """Load persisted momentum state on first access."""
    global _STATE_LOADED
    if _STATE_LOADED:
        return
    _STATE_LOADED = True
    try:
        from pulse.core.brain_utils import load_json_dict
        data = load_json_dict(_state_file())
        if data.get("consecutive_successes"):
            _CONSECUTIVE_SUCCESSES.update(data["consecutive_successes"])
    except Exception as e:
        log.debug("_load_state failed: %s", e)

def _save_state() -> None:
    """Persist momentum state after every update."""
    try:
        from pulse.core.brain_utils import atomic_update_json
        def _updater(_data):
            return {"consecutive_successes": dict(_CONSECUTIVE_SUCCESSES)}
        atomic_update_json(_state_file(), _updater, default={})
    except Exception as e:
        log.debug("_save_state failed: %s", e)
_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:
        _load_state()
        key = domain or "general"
        if success:
            streak_before = _CONSECUTIVE_SUCCESSES.get(key, 0)
            _CONSECUTIVE_SUCCESSES[key] = streak_before + 1
            streak = _CONSECUTIVE_SUCCESSES[key]
            _save_state()
            # 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
            _save_state()
            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:
        _load_state()
        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:
        _load_state()
        _CONSECUTIVE_SUCCESSES[domain or "general"] = 0
        _save_state()
    except Exception as _exc:
        pass  # QR14: silent — consider log.debug("%s", _exc)
