"""
Multi-objective reward function for CompanionGuard-RL intervention policy.

v5 design — label-aligned constrained reward:

Primary signal: alignment with human-annotated a_recommend
  +3.0  exact match with a_recommend (main training signal)
  +1.0  adjacent action (|action - a_recommend| == 1, partial credit)

Safety hard constraints:
  -5.0  PASS on L3/L4 (high-risk miss — must be penalized heavily)
  -2.0  PASS on L2 (moderate-risk miss)
  -5.0  strong intervention (REWRITE/REJECT/CRISIS) on L0 (over-refusal on safe content)

Category-specific bonuses and penalties:
  +3.0  CRISIS on R1 with l_risk >= 3 (crisis escalation for self-harm/suicide)
  +2.0  REJECT on R6/R7 with l_risk >= 3 (violence / minor inappropriate intimacy)
  +1.5  REWRITE on R3/R4/R10 with l_risk >= 2 (companion-specific adaptive rewriting)
  -3.0  CRISIS on samples that are neither R1 nor l_risk==4 (crisis misuse penalty)

Category index reference (PRIMARY_CATEGORY_LIST order):
  0=R1 Self-harm/Suicide  1=R2 Mental Health  2=R3 Dependency/Manipulation
  3=R4 Isolation          4=R5 Eating Disorder 5=R6 Violence
  6=R7 Minor Intimacy     7=R8 Privacy         8=R9 Decision Misguidance
  9=R10 Persona Immersion
"""

import torch
from typing import Union, Optional

ACTION_PASS    = 0
ACTION_WARN    = 1
ACTION_REWRITE = 2
ACTION_REJECT  = 3
ACTION_CRISIS  = 4

# Category indices (matches PRIMARY_CATEGORY_LIST order)
_R1_IDX  = 0   # Self-harm & Suicide Risk → CRISIS
_R3_IDX  = 2   # Emotional Dependency & Manipulation → REWRITE
_R4_IDX  = 3   # Real-world Support Isolation → REWRITE
_R6_IDX  = 5   # Violence & Dangerous Behavior → REJECT
_R7_IDX  = 6   # Minor Inappropriate Intimacy → REJECT
_R10_IDX = 9   # Persona Immersion Failure → REWRITE

# Reward weights
EXACT_ALIGN_BONUS     = 3.0   # action matches a_recommend exactly
ADJACENT_ALIGN_BONUS  = 1.0   # action is adjacent to a_recommend

PASS_HIGH_PENALTY     = -5.0  # PASS on L3/L4
PASS_MID_PENALTY      = -2.0  # PASS on L2
OVER_REFUSAL_PENALTY  = -5.0  # REWRITE/REJECT/CRISIS on L0

CRISIS_R1_BONUS       = 3.0   # CRISIS on R1 with l_risk >= 3
REJECT_R6R7_BONUS     = 2.0   # REJECT on R6/R7 with l_risk >= 3
REWRITE_COMPANION_BONUS = 1.5  # REWRITE on R3/R4/R10 with l_risk >= 2
CRISIS_MISUSE_PENALTY = -3.0  # CRISIS on non-R1 and l_risk < 4


def compute_reward(
    action: Union[int, torch.Tensor],
    y_risk: Union[int, torch.Tensor],
    l_risk: Union[int, torch.Tensor],
    c_primary_idx: Union[int, torch.Tensor],
    a_recommend: Optional[Union[int, torch.Tensor]] = None,
    weights: Optional[dict] = None,
) -> float:
    """
    Compute scalar reward for a single (state, action) step.

    Args:
        action:        intervention action taken (0–4)
        y_risk:        ground-truth binary risk label (0 or 1)
        l_risk:        ground-truth risk level (0–4)
        c_primary_idx: ground-truth primary category index (0=R1, …, 9=R10)
        a_recommend:   recommended action from annotation (primary alignment signal)
        weights:       reserved for future weight scaling (unused in v5)
    """
    if isinstance(action,        torch.Tensor): action        = int(action.item())
    if isinstance(y_risk,        torch.Tensor): y_risk        = int(y_risk.item())
    if isinstance(l_risk,        torch.Tensor): l_risk        = int(l_risk.item())
    if isinstance(c_primary_idx, torch.Tensor): c_primary_idx = int(c_primary_idx.item())
    if isinstance(a_recommend,   torch.Tensor): a_recommend   = int(a_recommend.item())

    l_risk = max(0, min(4, l_risk))
    action = max(0, min(4, action))

    reward = 0.0

    # 1. Primary signal: alignment with a_recommend
    if a_recommend is not None:
        if action == a_recommend:
            reward += EXACT_ALIGN_BONUS
        elif abs(action - a_recommend) == 1:
            reward += ADJACENT_ALIGN_BONUS

    # 2. Safety hard constraints (applied regardless of a_recommend)
    if action == ACTION_PASS:
        if l_risk >= 3:
            reward += PASS_HIGH_PENALTY
        elif l_risk == 2:
            reward += PASS_MID_PENALTY

    if l_risk == 0 and action >= ACTION_REWRITE:
        reward += OVER_REFUSAL_PENALTY

    # 3. Category-specific bonuses / penalties
    if action == ACTION_CRISIS:
        if c_primary_idx == _R1_IDX and l_risk >= 3:
            reward += CRISIS_R1_BONUS          # correct crisis escalation for self-harm
        if c_primary_idx != _R1_IDX and l_risk < 4:
            reward += CRISIS_MISUSE_PENALTY    # CRISIS misuse on non-crisis content

    if action == ACTION_REJECT:
        if c_primary_idx in (_R6_IDX, _R7_IDX) and l_risk >= 3:
            reward += REJECT_R6R7_BONUS        # violence / minor intimacy REJECT

    if action == ACTION_REWRITE:
        if c_primary_idx in (_R3_IDX, _R4_IDX, _R10_IDX) and l_risk >= 2:
            reward += REWRITE_COMPANION_BONUS  # companion-specific adaptive rewriting

    return reward


def compute_batch_reward(
    actions: torch.Tensor,
    y_risk: torch.Tensor,
    l_risk: torch.Tensor,
    c_primary_idx: torch.Tensor,
    a_recommend: Optional[torch.Tensor] = None,
    weights: Optional[dict] = None,
) -> torch.Tensor:
    """Vectorized batch reward computation."""
    rewards = torch.zeros(len(actions))
    for i in range(len(actions)):
        rec = a_recommend[i] if a_recommend is not None else None
        rewards[i] = compute_reward(
            actions[i], y_risk[i], l_risk[i], c_primary_idx[i],
            a_recommend=rec, weights=weights,
        )
    return rewards


# ── Backward-compatible alias ─────────────────────────────────────────────────
DEFAULT_WEIGHTS = {
    "w1": 2.0, "w2": 3.0, "w3": 4.0, "w4": 1.5, "w5": 0.5,
}
