env/grader.py

"""
SYNAPSE-X grader.

Score = 0.5 * completion_rate + 0.3 * efficiency + 0.2 * reward_score
"""

from typing import Any, Callable

from env.environment import SynapseXEnvironment
from env.models import Action, ActionPayload, GradeResult, Observation


EASY_TASK_CONFIG = [
    {
        "id": 0,
        "name": "Email Triage",
        "priority": 0.6,
        "risk": 0.1,
        "uncertainty": 0.05,
        "deadline": 18.0,
        "resources_required": 0.1,
    },
    {
        "id": 1,
        "name": "Report Drafting",
        "priority": 0.5,
        "risk": 0.15,
        "uncertainty": 0.1,
        "deadline": 16.0,
        "resources_required": 0.1,
    },
    {
        "id": 2,
        "name": "Data Backup",
        "priority": 0.4,
        "risk": 0.05,
        "uncertainty": 0.05,
        "deadline": 20.0,
        "resources_required": 0.05,
    },
]

MEDIUM_TASK_CONFIG = [
    {
        "id": 0,
        "name": "Code Review",
        "priority": 0.75,
        "risk": 0.26,
        "uncertainty": 0.28,
        "deadline": 11.0,
        "resources_required": 0.2,
    },
    {
        "id": 1,
        "name": "Customer Escalation",
        "priority": 0.85,
        "risk": 0.46,
        "uncertainty": 0.32,
        "deadline": 10.0,
        "resources_required": 0.25,
    },
    {
        "id": 2,
        "name": "Infrastructure Patch",
        "priority": 0.7,
        "risk": 0.42,
        "uncertainty": 0.35,
        "deadline": 10.0,
        "resources_required": 0.25,
    },
]

HARD_TASK_CONFIG = [
    {
        "id": 0,
        "name": "Live Incident Response",
        "priority": 1.0,
        "risk": 0.46,
        "uncertainty": 0.55,
        "deadline": 7.0,
        "resources_required": 0.2,
        "dependencies": [],
    },
    {
        "id": 1,
        "name": "Multi-system Migration",
        "priority": 0.9,
        "risk": 0.25,
        "uncertainty": 0.5,
        "deadline": 9.0,
        "resources_required": 0.3,
        "dependencies": [0],
    },
    {
        "id": 2,
        "name": "Regulatory Compliance Audit",
        "priority": 0.95,
        "risk": 0.54,
        "uncertainty": 0.60,
        "deadline": 10.0,
        "resources_required": 0.25,
        "dependencies": [0, 1],
    },
]

TRIAGE_TASK_CONFIG = [
    {
        "id": 0,
        "name": "Queue Warmup Review",
        "priority": 0.4,
        "risk": 0.08,
        "uncertainty": 0.08,
        "deadline": 5.0,
        "resources_required": 0.05,
        "release_time": 0,
    },
    {
        "id": 1,
        "name": "Urgent Abuse Escalation",
        "priority": 0.95,
        "risk": 0.18,
        "uncertainty": 0.12,
        "deadline": 4.0,
        "resources_required": 0.15,
        "release_time": 1,
    },
    {
        "id": 2,
        "name": "Coordinated Spam Wave",
        "priority": 0.72,
        "risk": 0.22,
        "uncertainty": 0.15,
        "deadline": 5.0,
        "resources_required": 0.1,
        "release_time": 1,
    },
    {
        "id": 3,
        "name": "Graphic Content Review",
        "priority": 0.88,
        "risk": 0.12,
        "uncertainty": 0.1,
        "deadline": 4.0,
        "resources_required": 0.15,
        "release_time": 2,
    },
    {
        "id": 4,
        "name": "VIP Appeal Audit",
        "priority": 0.8,
        "risk": 0.2,
        "uncertainty": 0.14,
        "deadline": 6.0,
        "resources_required": 0.1,
        "release_time": 2,
    },
]

TASK_REGISTRY: dict[str, list[dict[str, Any]]] = {
    "easy": EASY_TASK_CONFIG,
    "medium": MEDIUM_TASK_CONFIG,
    "hard": HARD_TASK_CONFIG,
    "triage": TRIAGE_TASK_CONFIG,
}

REWARD_BOUNDS = {
    "easy": {"min": -10.0, "max": 8.0},
    "medium": {"min": -6.0, "max": 7.0},
    "hard": {"min": -10.0, "max": 7.0},
    "triage": {"min": -8.0, "max": 8.0},
}

TASK_SEEDS = {"easy": 42, "medium": 42, "hard": 19}
TASK_SEEDS["triage"] = 31
TASK_SEED_VARIANTS = {
    "easy": [42, 43, 44],
    "medium": [42, 43, 44],
    "hard": [19, 20, 21],
    "triage": [31, 32, 33],
}


def _clip_unit_interval(value: float) -> float:
    return max(0.0, min(1.0, value))


def _build_env(task_name: str, seed: int) -> SynapseXEnvironment:
    if task_name not in TASK_REGISTRY:
        raise ValueError(f"Unknown task: {task_name}. Choose from {list(TASK_REGISTRY.keys())}")
    return SynapseXEnvironment(task_config=TASK_REGISTRY[task_name], seed=seed)


def _grade_for_seed(task_name: str, agent_actions: list[ActionPayload], seed: int) -> GradeResult:
    env = _build_env(task_name, seed)
    env.reset()
    total_reward = 0.0

    for raw_action in agent_actions:
        if env.done:
            break
        action = Action(**raw_action)
        result = env.step(action)
        total_reward += result.reward

    final_state = env.state()
    tasks = final_state["tasks"]

    completed = sum(1 for task in tasks if task["completed"])
    failed = sum(1 for task in tasks if task["failed"])
    total = len(tasks)

    completion_rate = completed / max(total, 1)
    efficiency = completed / max(completed + failed, 1)

    bounds = REWARD_BOUNDS[task_name]
    reward_score = (total_reward - bounds["min"]) / (bounds["max"] - bounds["min"])
    reward_score = _clip_unit_interval(reward_score)

    score = 0.5 * completion_rate + 0.3 * efficiency + 0.2 * reward_score
    score = round(_clip_unit_interval(score), 4)

    return GradeResult(
        score=score,
        completion_rate=round(completion_rate, 4),
        efficiency=round(efficiency, 4),
        reward_score=round(reward_score, 4),
        details={
            "task": task_name,
            "total_reward": round(total_reward, 4),
            "steps_taken": final_state["time"],
            "completed": completed,
            "failed": failed,
            "total_tasks": total,
        },
    )


def grade(task_name: str, agent_actions: list[ActionPayload]) -> GradeResult:
    return _grade_for_seed(task_name, agent_actions, TASK_SEEDS[task_name])


def grade_with_variance(
    task_name: str,
    action_generator: Callable[[str, Observation, SynapseXEnvironment], list[ActionPayload]],
    seeds: list[int] | None = None,
) -> GradeResult:
    selected_seeds = seeds or TASK_SEED_VARIANTS[task_name]
    seed_results: list[GradeResult] = []

    for seed in selected_seeds:
        env = _build_env(task_name, seed)
        obs = env.reset()
        actions = action_generator(task_name, obs, env)
        seed_results.append(_grade_for_seed(task_name, actions, seed))

    aggregate = {
        "score": round(sum(result.score for result in seed_results) / len(seed_results), 4),
        "completion_rate": round(sum(result.completion_rate for result in seed_results) / len(seed_results), 4),
        "efficiency": round(sum(result.efficiency for result in seed_results) / len(seed_results), 4),
        "reward_score": round(sum(result.reward_score for result in seed_results) / len(seed_results), 4),
    }
    return GradeResult(
        score=aggregate["score"],
        completion_rate=aggregate["completion_rate"],
        efficiency=aggregate["efficiency"],
        reward_score=aggregate["reward_score"],
        details={
            "task": task_name,
            "mode": "variance_mean",
            "seeds": selected_seeds,
            "seed_scores": [result.score for result in seed_results],
        },
    )


def run_all_tasks(
    agent_fn: Callable[[str, Observation, SynapseXEnvironment], list[ActionPayload]],
) -> dict[str, GradeResult]:
    results: dict[str, GradeResult] = {}
    for task_name in TASK_REGISTRY:
        env = _build_env(task_name, TASK_SEEDS[task_name])
        obs = env.reset()
        actions = agent_fn(task_name, obs, env)
        results[task_name] = grade(task_name, actions)
    return results