grid_env/tasks.py

"""
Task definitions for the warehouse fulfillment environment.
"""

from __future__ import annotations

from typing import Dict, List

from .models import BinState, OrderLine, TaskDefinition


GRID_SIZE = (7, 7)


def _base_bins() -> List[BinState]:
    return [
        BinState(bin_id="A1", position=(2, 1), sku="thermometer", quantity=2),
        BinState(bin_id="A2", position=(2, 3), sku="cough_syrup", quantity=2),
        BinState(bin_id="B1", position=(4, 1), sku="bandage_kit", quantity=2),
        BinState(bin_id="B2", position=(4, 3), sku="pain_relief", quantity=2),
        BinState(bin_id="C1", position=(2, 5), sku="gloves", quantity=2),
    ]


TASKS: Dict[str, TaskDefinition] = {
    "easy_single_pick": TaskDefinition(
        task_id="easy_single_pick",
        difficulty="easy",
        title="Single urgent order",
        description=(
            "Fulfill a same-hour pharmacy order by retrieving one thermometer "
            "from the correct bin and packing it at the station."
        ),
        max_steps=40,
        battery_capacity=36,
        low_battery_threshold=8,
        agent_start=(1, 1),
        agent_heading="E",
        dock_position=(1, 1),
        pack_station_position=(5, 5),
        charger_position=(1, 5),
        bins=_base_bins(),
        order=[OrderLine(sku="thermometer", quantity=1)],
        required_scans=["A1"],
        rubric_criteria=[
            {
                "name": "completion",
                "description": "Order fully packed.",
                "check": "param_at_least:state.completion_ratio=1.0",
            },
            {
                "name": "pick_item",
                "description": "Picked at least one item.",
                "check": "tool_used:pick_item",
            },
            {
                "name": "pack_item",
                "description": "Packed the order at the station.",
                "check": "tool_used:pack_item",
            },
            {
                "name": "no_wrong_picks",
                "description": "No incorrect picks.",
                "check": "param_at_most:state.wrong_picks=0",
            },
            {
                "name": "no_invalid_actions",
                "description": "No invalid actions.",
                "check": "param_at_most:state.invalid_actions=0",
            },
        ],
    ),
    "medium_multi_item": TaskDefinition(
        task_id="medium_multi_item",
        difficulty="medium",
        title="Two-line prescription basket",
        description=(
            "Fulfill a two-line order by scanning the relevant bins before each "
            "pick, then pack one cough syrup and one pain relief unit."
        ),
        max_steps=60,
        battery_capacity=34,
        low_battery_threshold=8,
        agent_start=(1, 1),
        agent_heading="E",
        dock_position=(1, 1),
        pack_station_position=(5, 5),
        charger_position=(1, 5),
        bins=_base_bins(),
        order=[
            OrderLine(sku="cough_syrup", quantity=1),
            OrderLine(sku="pain_relief", quantity=1),
        ],
        required_scans=["A2", "B2"],
        rubric_criteria=[
            {
                "name": "completion",
                "description": "All items packed.",
                "check": "param_at_least:state.completion_ratio=1.0",
            },
            {
                "name": "scans",
                "description": "Scanned required bins.",
                "check": "param_at_least:state.correct_scans=2",
            },
            {
                "name": "pack_item",
                "description": "Packed items at the station.",
                "check": "tool_used:pack_item",
            },
            {
                "name": "no_wrong_picks",
                "description": "No incorrect picks.",
                "check": "param_at_most:state.wrong_picks=0",
            },
            {
                "name": "no_wrong_scans",
                "description": "No incorrect scans.",
                "check": "param_at_most:state.wrong_scans=0",
            },
            {
                "name": "few_invalid_actions",
                "description": "At most one invalid action.",
                "check": "param_at_most:state.invalid_actions=1",
            },
        ],
    ),
    "hard_restock_priority": TaskDefinition(
        task_id="hard_restock_priority",
        difficulty="hard",
        title="Priority basket with battery management",
        description=(
            "Fulfill a three-line urgent order while managing battery reserve. "
            "The agent must scan each target bin, recharge when needed, and "
            "pack a thermometer, gloves, and bandage kit."
        ),
        max_steps=85,
        battery_capacity=24,
        low_battery_threshold=6,
        agent_start=(1, 1),
        agent_heading="E",
        dock_position=(1, 1),
        pack_station_position=(5, 5),
        charger_position=(1, 5),
        bins=_base_bins(),
        order=[
            OrderLine(sku="thermometer", quantity=1),
            OrderLine(sku="gloves", quantity=1),
            OrderLine(sku="bandage_kit", quantity=1),
        ],
        required_scans=["A1", "C1", "B1"],
        rubric_criteria=[
            {
                "name": "completion",
                "description": "All items packed.",
                "check": "param_at_least:state.completion_ratio=1.0",
            },
            {
                "name": "scans",
                "description": "Scanned all required bins.",
                "check": "param_at_least:state.correct_scans=3",
            },
            {
                "name": "recharge",
                "description": "Recharged at least once.",
                "check": "tool_used:recharge",
            },
            {
                "name": "no_battery_depletion",
                "description": "Avoided battery depletion.",
                "check": "param_at_most:state.battery_depletion_events=0",
            },
            {
                "name": "no_wrong_picks",
                "description": "No incorrect picks.",
                "check": "param_at_most:state.wrong_picks=0",
            },
            {
                "name": "few_invalid_actions",
                "description": "At most one invalid action.",
                "check": "param_at_most:state.invalid_actions=1",
            },
        ],
    ),

    # -----------------------------------------------------------------------
    # Task 4: obstacle_course (medium) — obstacles block direct paths
    # -----------------------------------------------------------------------
    "obstacle_course": TaskDefinition(
        task_id="obstacle_course",
        difficulty="medium",
        title="Obstacle-filled aisle navigation",
        description=(
            "Fulfill a two-item order in a warehouse cluttered with fallen crates. "
            "Navigate around obstacles to reach bins, scan them, pick one thermometer "
            "and one bandage kit, then pack both at the station."
        ),
        max_steps=70,
        battery_capacity=40,
        low_battery_threshold=10,
        agent_start=(0, 0),
        agent_heading="E",
        dock_position=(0, 0),
        pack_station_position=(6, 6),
        charger_position=(0, 6),
        bins=[
            BinState(bin_id="A1", position=(2, 1), sku="thermometer", quantity=2),
            BinState(bin_id="B1", position=(4, 1), sku="bandage_kit", quantity=2),
            BinState(bin_id="A2", position=(2, 3), sku="cough_syrup", quantity=2),
            BinState(bin_id="B2", position=(4, 3), sku="pain_relief", quantity=2),
            BinState(bin_id="C1", position=(2, 5), sku="gloves", quantity=2),
        ],
        order=[
            OrderLine(sku="thermometer", quantity=1),
            OrderLine(sku="bandage_kit", quantity=1),
        ],
        required_scans=["A1", "B1"],
        obstacles=[(1, 2), (2, 2), (3, 2), (3, 4), (4, 4), (5, 4)],
        rubric_criteria=[
            {
                "name": "completion",
                "description": "All items packed.",
                "check": "param_at_least:state.completion_ratio=1.0",
            },
            {
                "name": "scans",
                "description": "Scanned both required bins.",
                "check": "param_at_least:state.correct_scans=2",
            },
            {
                "name": "pack_item",
                "description": "Packed items at the station.",
                "check": "tool_used:pack_item",
            },
            {
                "name": "no_obstacle_collisions",
                "description": "Avoided all obstacle collisions.",
                "check": "param_at_most:state.obstacle_collisions=0",
            },
            {
                "name": "no_wrong_picks",
                "description": "No incorrect picks.",
                "check": "param_at_most:state.wrong_picks=0",
            },
            {
                "name": "few_invalid_actions",
                "description": "At most two invalid actions.",
                "check": "param_at_most:state.invalid_actions=2",
            },
        ],
    ),

    # -----------------------------------------------------------------------
    # Task 5: heavy_lifting (hard) — items have weight, limited carry capacity
    # -----------------------------------------------------------------------
    "heavy_lifting": TaskDefinition(
        task_id="heavy_lifting",
        difficulty="hard",
        title="Heavy-item logistics with weight limits",
        description=(
            "Fulfill a three-item order where items vary in weight (1-4 units). "
            "The agent has a carry capacity of 3 and must choose pickup order wisely. "
            "Heavier items drain more battery while moving. Scan each bin, pick items "
            "that fit within your carry limit, and pack at the station. "
            "The heavy pain_relief (weight 4) cannot be carried — skip it!"
        ),
        max_steps=90,
        battery_capacity=32,
        low_battery_threshold=8,
        agent_start=(1, 1),
        agent_heading="E",
        dock_position=(1, 1),
        pack_station_position=(5, 5),
        charger_position=(1, 5),
        bins=[
            BinState(bin_id="A1", position=(2, 1), sku="thermometer", quantity=3, weight=1),
            BinState(bin_id="A2", position=(2, 3), sku="cough_syrup", quantity=2, weight=2),
            BinState(bin_id="B1", position=(4, 1), sku="bandage_kit", quantity=2, weight=3),
            BinState(bin_id="B2", position=(4, 3), sku="pain_relief", quantity=2, weight=4),
            BinState(bin_id="C1", position=(2, 5), sku="gloves", quantity=3, weight=1),
        ],
        order=[
            OrderLine(sku="thermometer", quantity=1),
            OrderLine(sku="cough_syrup", quantity=1),
            OrderLine(sku="bandage_kit", quantity=1),
        ],
        required_scans=["A1", "A2", "B1"],
        carry_capacity=3,
        rubric_criteria=[
            {
                "name": "completion",
                "description": "All items packed.",
                "check": "param_at_least:state.completion_ratio=1.0",
            },
            {
                "name": "scans",
                "description": "Scanned all three required bins.",
                "check": "param_at_least:state.correct_scans=3",
            },
            {
                "name": "recharge",
                "description": "Recharged at least once.",
                "check": "tool_used:recharge",
            },
            {
                "name": "no_overweight",
                "description": "Never tried to pick an overweight item.",
                "check": "param_at_most:state.overweight_attempts=0",
            },
            {
                "name": "no_battery_depletion",
                "description": "Avoided battery depletion.",
                "check": "param_at_most:state.battery_depletion_events=0",
            },
            {
                "name": "no_wrong_picks",
                "description": "No incorrect picks.",
                "check": "param_at_most:state.wrong_picks=0",
            },
            {
                "name": "few_invalid_actions",
                "description": "At most two invalid actions.",
                "check": "param_at_most:state.invalid_actions=2",
            },
        ],
    ),

    # -----------------------------------------------------------------------
    # Task 6: stamina_run (hard) — stamina drains on movement
    # -----------------------------------------------------------------------
    "stamina_run": TaskDefinition(
        task_id="stamina_run",
        difficulty="hard",
        title="Endurance run with stamina management",
        description=(
            "Fulfill a two-item order while managing stamina. Every move drains stamina; "
            "when stamina hits zero, movement costs double battery. Rest at the rest area "
            "to restore stamina. Pick one cough syrup and one gloves unit, scan bins, and "
            "pack at the station without running out of energy."
        ),
        max_steps=80,
        battery_capacity=36,
        low_battery_threshold=8,
        agent_start=(0, 0),
        agent_heading="E",
        dock_position=(0, 0),
        pack_station_position=(6, 6),
        charger_position=(0, 6),
        rest_position=(3, 3),
        stamina_capacity=12,
        stamina_move_cost=1,
        bins=[
            BinState(bin_id="A1", position=(2, 1), sku="thermometer", quantity=2),
            BinState(bin_id="A2", position=(2, 3), sku="cough_syrup", quantity=2),
            BinState(bin_id="B1", position=(4, 1), sku="bandage_kit", quantity=2),
            BinState(bin_id="B2", position=(4, 3), sku="pain_relief", quantity=2),
            BinState(bin_id="C1", position=(2, 5), sku="gloves", quantity=2),
        ],
        order=[
            OrderLine(sku="cough_syrup", quantity=1),
            OrderLine(sku="gloves", quantity=1),
        ],
        required_scans=["A2", "C1"],
        rubric_criteria=[
            {
                "name": "completion",
                "description": "All items packed.",
                "check": "param_at_least:state.completion_ratio=1.0",
            },
            {
                "name": "scans",
                "description": "Scanned both required bins.",
                "check": "param_at_least:state.correct_scans=2",
            },
            {
                "name": "rest_used",
                "description": "Used the rest area at least once.",
                "check": "tool_used:rest",
            },
            {
                "name": "no_stamina_depletion",
                "description": "Avoided complete stamina depletion.",
                "check": "param_at_most:state.stamina_depletion_events=0",
            },
            {
                "name": "no_battery_depletion",
                "description": "Avoided battery depletion.",
                "check": "param_at_most:state.battery_depletion_events=0",
            },
            {
                "name": "no_wrong_picks",
                "description": "No incorrect picks.",
                "check": "param_at_most:state.wrong_picks=0",
            },
        ],
    ),

    # -----------------------------------------------------------------------
    # Task 7: budget_run (expert) — money rewards and profit target
    # -----------------------------------------------------------------------
    "budget_run": TaskDefinition(
        task_id="budget_run",
        difficulty="expert",
        title="Profitable fulfillment under budget pressure",
        description=(
            "Fulfill orders for profit. Each item has a dollar value earned when correctly "
            "packed. Wrong packs lose half the item value. You must reach a profit target "
            "of $15.00 while completing the order. Pick high-value items efficiently: "
            "2 thermometers ($5 each) and 1 bandage kit ($8). Budget-aware decisions matter."
        ),
        max_steps=70,
        battery_capacity=30,
        low_battery_threshold=6,
        agent_start=(1, 1),
        agent_heading="E",
        dock_position=(1, 1),
        pack_station_position=(5, 5),
        charger_position=(1, 5),
        bins=[
            BinState(bin_id="A1", position=(2, 1), sku="thermometer", quantity=3, value=5.0),
            BinState(bin_id="A2", position=(2, 3), sku="cough_syrup", quantity=2, value=3.0),
            BinState(bin_id="B1", position=(4, 1), sku="bandage_kit", quantity=2, value=8.0),
            BinState(bin_id="B2", position=(4, 3), sku="pain_relief", quantity=2, value=4.0),
            BinState(bin_id="C1", position=(2, 5), sku="gloves", quantity=2, value=2.0),
        ],
        order=[
            OrderLine(sku="thermometer", quantity=2),
            OrderLine(sku="bandage_kit", quantity=1),
        ],
        required_scans=["A1", "B1"],
        profit_target=15.0,
        rubric_criteria=[
            {
                "name": "completion",
                "description": "All items packed.",
                "check": "param_at_least:state.completion_ratio=1.0",
            },
            {
                "name": "profit_target",
                "description": "Reached the profit target of $15.",
                "check": "param_at_least:state.money_earned=15.0",
            },
            {
                "name": "scans",
                "description": "Scanned required bins.",
                "check": "param_at_least:state.correct_scans=2",
            },
            {
                "name": "recharge",
                "description": "Recharged at least once.",
                "check": "tool_used:recharge",
            },
            {
                "name": "no_money_lost",
                "description": "No money lost from wrong packs.",
                "check": "param_at_most:state.money_lost=0.0",
            },
            {
                "name": "no_wrong_picks",
                "description": "No incorrect picks.",
                "check": "param_at_most:state.wrong_picks=0",
            },
            {
                "name": "few_invalid_actions",
                "description": "At most one invalid action.",
                "check": "param_at_most:state.invalid_actions=1",
            },
        ],
    ),

    # -----------------------------------------------------------------------
    # Task 8: gauntlet (expert) — all mechanics combined
    # -----------------------------------------------------------------------
    "gauntlet": TaskDefinition(
        task_id="gauntlet",
        difficulty="expert",
        title="The gauntlet: obstacles, weight, stamina, and profit",
        description=(
            "The ultimate warehouse challenge. Navigate a cluttered warehouse with obstacles, "
            "manage item weights (carry capacity 3), conserve stamina (rest when needed), "
            "earn money for packed items, and hit a $20 profit target. Fulfill a four-item "
            "order: 1 thermometer ($5, wt 1), 1 cough syrup ($6, wt 2), 1 bandage kit ($8, wt 3), "
            "and 1 gloves ($4, wt 1). Recharge battery, rest for stamina, avoid obstacles, "
            "and finish profitable."
        ),
        max_steps=120,
        battery_capacity=28,
        low_battery_threshold=7,
        agent_start=(0, 0),
        agent_heading="S",
        dock_position=(0, 0),
        pack_station_position=(6, 6),
        charger_position=(6, 0),
        rest_position=(0, 6),
        stamina_capacity=10,
        stamina_move_cost=1,
        carry_capacity=3,
        profit_target=20.0,
        obstacles=[(1, 1), (3, 1), (5, 3), (3, 3), (1, 5), (5, 5)],
        bins=[
            BinState(bin_id="A1", position=(2, 0), sku="thermometer", quantity=3, weight=1, value=5.0),
            BinState(bin_id="A2", position=(2, 4), sku="cough_syrup", quantity=2, weight=2, value=6.0),
            BinState(bin_id="B1", position=(4, 0), sku="bandage_kit", quantity=2, weight=3, value=8.0),
            BinState(bin_id="B2", position=(4, 4), sku="pain_relief", quantity=2, weight=4, value=4.0),
            BinState(bin_id="C1", position=(4, 2), sku="gloves", quantity=3, weight=1, value=4.0),
        ],
        order=[
            OrderLine(sku="thermometer", quantity=1),
            OrderLine(sku="cough_syrup", quantity=1),
            OrderLine(sku="bandage_kit", quantity=1),
            OrderLine(sku="gloves", quantity=1),
        ],
        required_scans=["A1", "A2", "B1", "C1"],
        rubric_criteria=[
            {
                "name": "completion",
                "description": "All four items packed.",
                "check": "param_at_least:state.completion_ratio=1.0",
            },
            {
                "name": "profit_target",
                "description": "Reached $20 profit target.",
                "check": "param_at_least:state.money_earned=20.0",
            },
            {
                "name": "scans",
                "description": "Scanned all four required bins.",
                "check": "param_at_least:state.correct_scans=4",
            },
            {
                "name": "recharge",
                "description": "Recharged at least once.",
                "check": "tool_used:recharge",
            },
            {
                "name": "rest_used",
                "description": "Used the rest area at least once.",
                "check": "tool_used:rest",
            },
            {
                "name": "no_obstacle_collisions",
                "description": "Avoided all obstacle collisions.",
                "check": "param_at_most:state.obstacle_collisions=0",
            },
            {
                "name": "no_overweight",
                "description": "Never tried to pick an overweight item.",
                "check": "param_at_most:state.overweight_attempts=0",
            },
            {
                "name": "no_battery_depletion",
                "description": "Avoided battery depletion.",
                "check": "param_at_most:state.battery_depletion_events=0",
            },
            {
                "name": "no_stamina_depletion",
                "description": "Avoided stamina depletion.",
                "check": "param_at_most:state.stamina_depletion_events=0",
            },
            {
                "name": "no_wrong_picks",
                "description": "No incorrect picks.",
                "check": "param_at_most:state.wrong_picks=0",
            },
            {
                "name": "few_invalid_actions",
                "description": "At most two invalid actions.",
                "check": "param_at_most:state.invalid_actions=2",
            },
        ],
    ),
}


def get_task(task_id: str) -> TaskDefinition:
    if task_id not in TASKS:
        raise KeyError(f"Unknown task_id: {task_id}")
    return TASKS[task_id]