env/reward.py

"""
Per-step continuous reward function.
Rewards are given each step based on the quality of the action taken.
"""

from typing import List, Optional
from .models import Action, Email, Goal, EnvState


def compute_reward(
    action: Action,
    email: Email,
    goals: List[Goal],
    state: EnvState,
) -> float:
    """
    Compute per-step reward for taking `action` on `email`.

    Component breakdown (for interpretability):
    ┌─ POSITIVE ──────────────────────────────────────────────────────────┐
    │ goal_completion    +0.30 * (priority/5)  Completing a defined goal  │
    │ action_match       +0.10                 Correct action type chosen  │
    │ urgency_bonus      +0.15 * urgency       Acting fast on urgent items │
    │ delegation_bonus   +0.05                 Smart delegation of P1-P2   │
    ├─ NEGATIVE ──────────────────────────────────────────────────────────┤
    │ step_cost          -0.02                 Baseline cost per step      │
    │ ignore_critical    -0.25                 Ignoring P4-P5 emails       │
    │ wasted_step        -0.10                 Replying to P1 over P4+     │
    └─────────────────────────────────────────────────────────────────────┘

    Returns: reward clamped to [-0.3, 0.5]
    """
    reward = 0.0

    # --- Per-step cost (encourages efficiency) ---
    reward -= 0.02

    # --- Find if this action completes a goal ---
    completed_goal: Optional[Goal] = None
    for goal in goals:
        if goal.target_email_id == email.id and not goal.completed:
            if goal.required_action == action.type:
                completed_goal = goal
                break

    # --- Goal completion reward ---
    if completed_goal is not None:
        reward += 0.30 * (completed_goal.priority / 5.0)

    # --- Urgency bonus: act fast on high-urgency items ---
    if action.type != "ignore" and email.urgency > 0.5:
        reward += 0.15 * email.urgency

    # --- Efficiency: delegating low-priority items is smart ---
    if action.type == "delegate" and email.priority <= 2:
        reward += 0.05

    # --- Appropriate action bonus ---
    if completed_goal is not None:
        reward += 0.10

    # --- Penalty: ignoring critical emails ---
    if action.type == "ignore" and email.priority >= 4:
        reward -= 0.25

    # --- Penalty: wasting steps on low-priority when high-priority unhandled ---
    if email.priority <= 2:
        unhandled_critical = [
            e for e in state.inbox
            if not e.handled and e.priority >= 4 and e.id != email.id
        ]
        if unhandled_critical and action.type == "reply":
            reward -= 0.10

    return max(-0.3, min(0.5, reward))