rewards.py

"""
FairRecovery — Reward Engine (RLVR).

Computes dense, verifiable, formula-based rewards — no learned reward model.
All reward components are returned separately for transparency and logging.

R_total = w_exec * R_exec + w_fair * R_fair + w_safe * R_safe

  R_exec — average service increase from allocations (execution quality)
  R_fair — negative disparity between vulnerable vs non-vulnerable group service
  R_safe — negative penalty for constraint violations

Anti-reward-hacking properties:
  • Vulnerable-zone ignore penalty: agent can't just ignore high-vulnerability zones
  • Budget overflow blocked: allocations exceeding budget are voided, not scaled
  • Per-component caps: no single component can dominate to mask failures elsewhere
"""

from __future__ import annotations

import structlog
from dataclasses import dataclass, field
from typing import List, Optional, Tuple

from .constants import (
    GRADER_SCORE_MAX,
    GRADER_SCORE_MIN,
    MAX_DAYS,
    PENALTY_IGNORE_VULNERABLE,
    PENALTY_REPEATED_ACTION,
    REWARD_WEIGHTS,
    VULNERABILITY_THRESHOLD,
)
from .state import CityState, ZoneState
from .tasks import ScenarioConfig

logger = structlog.get_logger(__name__)


# ──────────────────────────────────────────────────────────────────────────────
# Component reward functions (pure, testable)
# ──────────────────────────────────────────────────────────────────────────────

def compute_exec_reward(prev_services: List[float], zones: List[ZoneState]) -> float:
    """
    Mean service improvement this day.

    Positive when allocations increased service levels.
    Zero when no allocations were made or budget was exhausted.
    """
    if not zones:
        return 0.0
    improvements = [
        z.service - prev
        for z, prev in zip(zones, prev_services)
    ]
    return float(sum(improvements) / len(improvements))


def compute_fairness_reward(zones: List[ZoneState]) -> float:
    """
    Negative disparity between vulnerable and non-vulnerable group service.

    Returns a value ≤ 0:
      0.0   = perfect parity (vulnerable zones served as well as others)
      -1.0  = maximum disparity (vulnerable zones have 0 service, others have 1)

    Anti-gaming: measures group averages, not individual zones,
    so agent can't game it by over-serving one token vulnerable zone.
    """
    vuln   = [z for z in zones if z.vulnerable_ratio >= VULNERABILITY_THRESHOLD]
    normal = [z for z in zones if z.vulnerable_ratio <  VULNERABILITY_THRESHOLD]

    if not vuln or not normal:
        # Homogeneous population — use service variance as proxy
        services = [z.service for z in zones]
        if len(services) < 2:
            return 0.0
        mean_svc = sum(services) / len(services)
        variance = sum((s - mean_svc) ** 2 for s in services) / len(services)
        return float(-variance)

    avg_vuln   = sum(z.service for z in vuln)   / len(vuln)
    avg_normal = sum(z.service for z in normal) / len(normal)
    disparity  = avg_normal - avg_vuln    # positive = normal zones better off
    return float(-disparity)              # penalise if vulnerable lag behind


def compute_safety_reward(violations: List[str]) -> float:
    """
    Penalty per safety violation, capped to avoid dominating the total.
    -0.1 per violation, max -0.5 total.
    """
    penalty = min(0.5, len(violations) * 0.1)
    return float(-penalty)


def compute_analysis_reward(
    chosen_zones: List[int],
    zones: List[ZoneState],
) -> float:
    """
    Partial reward for correctly identifying the most critical zones.

    Correct = chose at least one zone from the top-k by (damage × vulnerable_ratio).
    Returns proportion of top-k correctly identified.
    """
    if not zones or not chosen_zones:
        return 0.0

    k = max(1, len(zones) // 2)
    ranked = sorted(
        range(len(zones)),
        key=lambda i: zones[i].damage * zones[i].vulnerable_ratio,
        reverse=True,
    )
    top_k   = set(ranked[:k])
    chosen  = set(chosen_zones)
    overlap = len(top_k & chosen)
    return float(overlap / k)


@dataclass
class RewardComponents:
    """Named reward components for a single step — logged and returned to training."""
    R_exec: float = 0.0
    R_fair: float = 0.0
    R_safe: float = 0.0
    R_analysis: float = 0.0
    R_total: float = 0.0
    violations: List[str] = field(default_factory=list)
    feedback: str = ""

    def to_dict(self) -> dict:
        return {
            "R_exec":     round(self.R_exec, 4),
            "R_fair":     round(self.R_fair, 4),
            "R_safe":     round(self.R_safe, 4),
            "R_analysis": round(self.R_analysis, 4),
            "R_total":    round(self.R_total, 4),
            "violations": self.violations,
        }


# ──────────────────────────────────────────────────────────────────────────────
# Stateful Reward Engine
# ──────────────────────────────────────────────────────────────────────────────

class RewardEngine:
    """
    Stateful reward calculator for a single episode.

    Tracks cumulative reward and step count.
    All individual computations delegate to pure functions above.
    """

    def __init__(self, task: ScenarioConfig) -> None:
        self._task = task
        self._cumulative_reward: float = 0.0
        self._step_count: int = 0
        self._action_history: List[str] = []
        self._vulnerable_ignored_days: int = 0

    @property
    def cumulative_reward(self) -> float:
        return self._cumulative_reward

    # ── analysis step reward ─────────────────────────────────────────────────

    def compute_analysis_step(
        self,
        chosen_zones: List[int],
        city: CityState,
    ) -> RewardComponents:
        """Reward for analysis quality. Dense partial-progress signal."""
        self._step_count += 1
        R_analysis = compute_analysis_reward(chosen_zones, city.zones)
        R_total = 0.1 * R_analysis  # small reward — analysis is not execution

        self._cumulative_reward += R_total
        components = RewardComponents(
            R_analysis=R_analysis,
            R_total=R_total,
            feedback=(
                f"Analysis reward: {R_total:.3f} "
                f"(identified {int(R_analysis * max(1, len(city.zones)//2))}"
                f"/{max(1, len(city.zones)//2)} critical zones correctly)."
            ),
        )
        self._log("analysis", components)
        return components

    # ── execute step reward ───────────────────────────────────────────────────

    def compute_execute_step(
        self,
        city: CityState,
        violations: List[str],
    ) -> RewardComponents:
        """
        Main dense reward after execute step.

        Also checks if vulnerable zones were consistently ignored —
        a key anti-reward-hacking signal.
        """
        self._step_count += 1

        # Check if high-vulnerability zones were ignored this day
        vuln_zone_ids = {z.zone_id for z in city.zones if z.is_vulnerable}
        allocated_zone_ids = {
            a.get("zone") for a in (city.pending_allocations or [])
        } if city.pending_allocations else set()

        # Check historically — did any vulnerable zone receive resources?
        if vuln_zone_ids:
            history_text = " ".join(city.history)
            zone_served = any(
                f"zone {zid}" in history_text.lower() or str(zid) in history_text
                for zid in vuln_zone_ids
            )
            if not zone_served and city.day > 1:
                self._vulnerable_ignored_days += 1
                if self._vulnerable_ignored_days >= 2:
                    violations.append(f"persistent_ignore_vulnerable:{vuln_zone_ids}")

        R_exec = compute_exec_reward(city.prev_services, city.zones)
        R_fair = compute_fairness_reward(city.zones)
        R_safe = compute_safety_reward(violations)

        w = REWARD_WEIGHTS
        R_total = (
            w["exec"] * R_exec +
            w["fair"] * R_fair +
            w["safe"] * R_safe
        )

        # Clamp to [-1, 1] per step
        R_total = float(max(-1.0, min(1.0, R_total)))

        self._cumulative_reward += R_total

        feedback_parts = [
            f"R_exec={R_exec:+.3f}",
            f"R_fair={R_fair:+.3f}",
            f"R_safe={R_safe:+.3f}",
            f"→ R_total={R_total:+.3f}",
        ]
        if violations:
            feedback_parts.append(f"Violations: {violations}.")

        components = RewardComponents(
            R_exec=R_exec,
            R_fair=R_fair,
            R_safe=R_safe,
            R_total=R_total,
            violations=violations,
            feedback=" | ".join(feedback_parts),
        )
        self._log("execute", components)
        return components

    # ── final episode reward ─────────────────────────────────────────────────

    def compute_submit_reward(self, city: CityState) -> RewardComponents:
        """
        Terminal reward when agent submits.

        Combines final service level and final fairness score
        into a normalised bonus. Efficient episodes (fewer steps) earn a bonus.
        """
        self._step_count += 1

        R_fair = compute_fairness_reward(city.zones)
        avg_svc = sum(z.service for z in city.zones) / max(1, len(city.zones))

        # Terminal bonus in [0, 1]:  0.5*avg_service + 0.5*(1 + R_fair) clamped
        terminal = 0.5 * avg_svc + 0.5 * (1.0 + R_fair)
        terminal = float(max(0.0, min(1.0, terminal)))

        self._cumulative_reward += terminal

        components = RewardComponents(
            R_fair=R_fair,
            R_exec=avg_svc,
            R_total=terminal,
            feedback=(
                f"Episode submitted. Terminal bonus={terminal:.3f} "
                f"(avg_service={avg_svc:.3f}, fairness={R_fair:.3f})."
            ),
        )
        self._log("submit", components)
        return components

    # ── grader score ─────────────────────────────────────────────────────────

    def get_final_grader_score(self, city: CityState) -> float:
        """
        Normalised score in (GRADER_SCORE_MIN, GRADER_SCORE_MAX).

        Based on: average service improvement + fairness improvement.
        Never exactly 0 or 1 (per OpenEnv grading conventions).
        """
        avg_svc = sum(z.service for z in city.zones) / max(1, len(city.zones))
        R_fair  = compute_fairness_reward(city.zones)

        # Normalise: avg_svc in [0,1], R_fair in [-1,0] → map to [0,1]
        normalised = 0.6 * avg_svc + 0.4 * (1.0 + R_fair)
        clamped    = max(GRADER_SCORE_MIN, min(GRADER_SCORE_MAX, normalised))
        return round(float(clamped), 4)

    # ── logging ───────────────────────────────────────────────────────────────

    def _log(self, step_type: str, components: RewardComponents) -> None:
        logger.info(
            "reward_computed",
            step_type=step_type,
            step=self._step_count,
            **components.to_dict(),
            cumulative=round(self._cumulative_reward, 4),
        )