src/grader.py

"""
Task graders for the Email Triage environment.

Each grader takes an episode history (list of StepRecords) and returns a
deterministic score strictly within (0.0, 1.0) — never exactly 0 or 1.

Three graders implement progressive difficulty:
  - grade_task_basic:   5 emails, 2 folders, accuracy-only scoring
  - grade_task_medium:  15 emails, 4 folders, weighted per-folder accuracy
  - grade_task_hard:    30 emails, 6 folders, multi-component with VIP/urgency
"""

from __future__ import annotations

import logging
from typing import Dict, List

from src.models import StepRecord

logger = logging.getLogger(__name__)

# Minimum accuracy thresholds
MEDIUM_THRESHOLD: float = 0.70
HARD_ACCURACY_THRESHOLD: float = 0.60

# Epsilon bounds to keep scores strictly in (0, 1)
_SCORE_MIN: float = 0.01
_SCORE_MAX: float = 0.99


def _clamp_score(score: float) -> float:
    """Clamp a score to the open interval (0, 1).

    The OpenEnv evaluation requires scores strictly between 0 and 1 —
    never exactly 0.0 or 1.0.
    """
    return round(max(_SCORE_MIN, min(_SCORE_MAX, score)), 4)


def grade_task_basic(episode_history: List[StepRecord]) -> float:
    """Grade the basic email sorting task.

    Task: Sort 5 emails into work vs. spam.
    Scoring: Pure accuracy (fraction of emails placed in correct folder).

    Args:
        episode_history: List of StepRecords from the episode.

    Returns:
        Score strictly in (0.0, 1.0).
    """
    if not episode_history:
        return _SCORE_MIN

    move_steps = _get_classification_steps(episode_history)
    if not move_steps:
        return _SCORE_MIN

    correct = 0.0
    for step in move_steps:
        c = step.reward.components.get("correctness", 0.0)
        if c >= 0.9:
            correct += 1.0
        elif c >= 0.4:
            correct += 0.5  # partial credit

    total_emails = max(len(move_steps), 5)
    score = correct / total_emails

    return _clamp_score(score)


def grade_task_medium(episode_history: List[StepRecord]) -> float:
    """Grade the multi-folder triage task.

    Task: Sort 15 emails into 4 folders (work, finance, meetings, spam).
    Scoring: Weighted per-folder accuracy, with a 70% threshold gate.

    Folder weights reflect business importance:
      - work:     0.35
      - finance:  0.30
      - meetings: 0.20
      - spam:     0.15

    Args:
        episode_history: List of StepRecords from the episode.

    Returns:
        Score strictly in (0.0, 1.0).
    """
    if not episode_history:
        return _SCORE_MIN

    move_steps = _get_classification_steps(episode_history)
    if not move_steps:
        return _SCORE_MIN

    folder_weights: Dict[str, float] = {
        "work": 0.35,
        "finance": 0.30,
        "meetings": 0.20,
        "spam": 0.15,
    }

    folder_correct: Dict[str, float] = {}
    folder_total: Dict[str, int] = {}

    for step in move_steps:
        gt_folder = step.email.ground_truth_folder
        if gt_folder not in folder_weights:
            gt_folder = "work"

        folder_total[gt_folder] = folder_total.get(gt_folder, 0) + 1

        c = step.reward.components.get("correctness", 0.0)
        if c >= 0.9:
            folder_correct[gt_folder] = folder_correct.get(gt_folder, 0.0) + 1.0
        elif c >= 0.4:
            folder_correct[gt_folder] = folder_correct.get(gt_folder, 0.0) + 0.5

    weighted_score = 0.0
    for folder, weight in folder_weights.items():
        total = folder_total.get(folder, 0)
        if total > 0:
            accuracy = folder_correct.get(folder, 0.0) / total
            weighted_score += weight * accuracy

    active_weight = sum(
        w for f, w in folder_weights.items() if folder_total.get(f, 0) > 0
    )
    if active_weight > 0:
        weighted_score = weighted_score / active_weight

    # Threshold gate: low accuracy gets a low (but non-zero) score
    overall_accuracy = _overall_accuracy(move_steps)
    if overall_accuracy < MEDIUM_THRESHOLD:
        return _clamp_score(weighted_score * 0.3)

    return _clamp_score(weighted_score)


def grade_task_hard(episode_history: List[StepRecord]) -> float:
    """Grade the advanced triage task with urgency and VIP handling.

    Task: Sort 30 emails with deadline awareness and VIP prioritization.
    Scoring: Multi-component.
      - 50% Overall accuracy (urgent emails weighted 2x)
      - 25% Efficiency (fewer steps = better)
      - 25% VIP handling (correctly classified VIP emails)

    Args:
        episode_history: List of StepRecords from the episode.

    Returns:
        Score strictly in (0.0, 1.0).
    """
    if not episode_history:
        return _SCORE_MIN

    move_steps = _get_classification_steps(episode_history)
    if not move_steps:
        return _SCORE_MIN

    # ── 1. Accuracy (50%) — urgent emails weighted 2x ────────────────────
    weighted_correct = 0.0
    weighted_total = 0.0

    for step in move_steps:
        weight = 2.0 if step.email.priority_flag else 1.0
        weighted_total += weight
        c = step.reward.components.get("correctness", 0.0)
        if c >= 0.9:
            weighted_correct += weight
        elif c >= 0.4:
            weighted_correct += weight * 0.5

    accuracy_score = weighted_correct / max(weighted_total, 1.0)

    # ── 2. Efficiency (25%) ──────────────────────────────────────────────
    total_steps = len(episode_history)
    ideal_steps = 30
    efficiency_score = max(0.0, 1.0 - max(0, total_steps - ideal_steps) / ideal_steps)

    # ── 3. VIP handling (25%) ────────────────────────────────────────────
    vip_steps = [s for s in move_steps if s.email.is_vip_sender]
    if vip_steps:
        vip_correct = sum(
            1.0 for s in vip_steps
            if s.reward.components.get("correctness", 0.0) >= 0.9
        )
        vip_partial = sum(
            0.5 for s in vip_steps
            if 0.4 <= s.reward.components.get("correctness", 0.0) < 0.9
        )
        vip_score = (vip_correct + vip_partial) / len(vip_steps)
    else:
        vip_score = 0.0

    # ── Combine ──────────────────────────────────────────────────────────
    final_score = (
        0.50 * accuracy_score
        + 0.25 * efficiency_score
        + 0.25 * vip_score
    )

    # Gate: low accuracy gets a penalized (but non-zero) score
    raw_accuracy = _overall_accuracy(move_steps)
    if raw_accuracy < HARD_ACCURACY_THRESHOLD:
        return _clamp_score(final_score * 0.3)

    return _clamp_score(final_score)


# ── Helper Functions ─────────────────────────────────────────────────────────

def _get_classification_steps(history: List[StepRecord]) -> List[StepRecord]:
    """Filter to steps that classify emails (move, mark_spam, delete)."""
    classification_actions = {"move", "mark_spam", "delete"}
    return [
        s for s in history
        if s.action.action_type in classification_actions
    ]


def _overall_accuracy(move_steps: List[StepRecord]) -> float:
    """Simple accuracy: fraction of correctly classified emails."""
    if not move_steps:
        return 0.0
    correct = sum(
        1.0 for s in move_steps
        if s.reward.components.get("correctness", 0.0) >= 0.9
    )
    return correct / len(move_steps)