tasks.py

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Dict, List

from .models import DeliveryWindow, PendingOrder


class Task(ABC):
    name: str
    num_platforms: int
    max_steps: int
    seed: int
    solar_zones: Dict[str, float]
    pending_orders: List[PendingOrder]
    delivery_windows: List[DeliveryWindow]
    description: str

    # Grading thresholds
    min_production_runs: int
    min_assemblies: int
    min_deliveries: int
    min_energy_final: float
    max_invalid_actions: int


class EasyTask(Task):
    name = "easy"
    num_platforms = 2
    max_steps = 60
    seed = 101
    description = "Basic production scheduling with 2 platforms and 4 delivery orders."
    solar_zones = {"zone_a": 0.8, "zone_b": 0.8}

    pending_orders = [
        PendingOrder(order_id=1, product_type="alloy_panel", requires_assembly=False),
        PendingOrder(order_id=2, product_type="circuit_board", requires_assembly=False),
        PendingOrder(order_id=3, product_type="fuel_cell", requires_assembly=False),
        PendingOrder(order_id=4, product_type="alloy_panel", requires_assembly=False),
    ]
    delivery_windows = [
        DeliveryWindow(order_id=1, product_type="alloy_panel", deadline=20, reward_value=8.0),
        DeliveryWindow(order_id=2, product_type="circuit_board", deadline=35, reward_value=8.0),
        DeliveryWindow(order_id=3, product_type="fuel_cell", deadline=50, reward_value=8.0),
        DeliveryWindow(order_id=4, product_type="alloy_panel", deadline=58, reward_value=8.0),
    ]

    min_production_runs = 3
    min_assemblies = 0
    min_deliveries = 2
    min_energy_final = 40.0
    max_invalid_actions = 3


class MediumTask(Task):
    name = "medium"
    num_platforms = 4
    max_steps = 120
    seed = 202
    description = "Multi-stage assembly and delivery balancing with 4 platforms and 10 orders."
    solar_zones = {
        "zone_a": 0.8,
        "zone_b": 0.6,
        "zone_c": 0.5,
        "zone_d": 0.7,
    }

    pending_orders = [
        PendingOrder(order_id=1, product_type="alloy_panel", requires_assembly=False),
        PendingOrder(order_id=2, product_type="circuit_board", requires_assembly=False),
        PendingOrder(order_id=3, product_type="solar_array", requires_assembly=True),
        PendingOrder(order_id=4, product_type="fuel_cell", requires_assembly=False),
        PendingOrder(order_id=5, product_type="thruster_module", requires_assembly=True),
        PendingOrder(order_id=6, product_type="alloy_panel", requires_assembly=False),
        PendingOrder(order_id=7, product_type="circuit_board", requires_assembly=False),
        PendingOrder(order_id=8, product_type="solar_array", requires_assembly=True),
        PendingOrder(order_id=9, product_type="fuel_cell", requires_assembly=False),
        PendingOrder(order_id=10, product_type="habitat_unit", requires_assembly=True),
    ]
    delivery_windows = [
        DeliveryWindow(order_id=1, product_type="alloy_panel", deadline=25, reward_value=8.0),
        DeliveryWindow(order_id=2, product_type="circuit_board", deadline=40, reward_value=8.0),
        DeliveryWindow(order_id=3, product_type="solar_array", deadline=55, reward_value=10.0),
        DeliveryWindow(order_id=4, product_type="fuel_cell", deadline=60, reward_value=8.0),
        DeliveryWindow(order_id=5, product_type="thruster_module", deadline=75, reward_value=10.0),
        DeliveryWindow(order_id=6, product_type="alloy_panel", deadline=80, reward_value=8.0),
        DeliveryWindow(order_id=7, product_type="circuit_board", deadline=90, reward_value=8.0),
        DeliveryWindow(order_id=8, product_type="solar_array", deadline=100, reward_value=10.0),
        DeliveryWindow(order_id=9, product_type="fuel_cell", deadline=110, reward_value=8.0),
        DeliveryWindow(order_id=10, product_type="habitat_unit", deadline=118, reward_value=12.0),
    ]

    min_production_runs = 8
    min_assemblies = 4
    min_deliveries = 6
    min_energy_final = 25.0
    max_invalid_actions = 8


class HardTask(Task):
    name = "hard"
    num_platforms = 6
    max_steps = 240
    seed = 303
    description = "Power pressure, multi-platform coordination, 20 orders, strict grading."
    solar_zones = {
        "zone_a": 0.9,
        "zone_b": 0.2,
        "zone_c": 0.6,
        "zone_d": 0.4,
        "zone_e": 0.7,
        "zone_f": 0.3,
    }

    pending_orders = [
        PendingOrder(order_id=i, product_type=pt, requires_assembly=asm)
        for i, (pt, asm) in enumerate(
            [
                ("alloy_panel", False),
                ("circuit_board", False),
                ("solar_array", True),
                ("fuel_cell", False),
                ("thruster_module", True),
                ("habitat_unit", True),
                ("alloy_panel", False),
                ("circuit_board", False),
                ("solar_array", True),
                ("fuel_cell", False),
                ("thruster_module", True),
                ("alloy_panel", False),
                ("circuit_board", False),
                ("habitat_unit", True),
                ("fuel_cell", False),
                ("solar_array", True),
                ("alloy_panel", False),
                ("thruster_module", True),
                ("circuit_board", False),
                ("habitat_unit", True),
            ],
            start=1,
        )
    ]
    delivery_windows = [
        DeliveryWindow(order_id=i, product_type=pt, deadline=dl, reward_value=rv)
        for i, pt, dl, rv in [
            (1, "alloy_panel", 20, 8.0),
            (2, "circuit_board", 35, 8.0),
            (3, "solar_array", 50, 10.0),
            (4, "fuel_cell", 60, 8.0),
            (5, "thruster_module", 70, 10.0),
            (6, "habitat_unit", 85, 12.0),
            (7, "alloy_panel", 90, 8.0),
            (8, "circuit_board", 100, 8.0),
            (9, "solar_array", 110, 10.0),
            (10, "fuel_cell", 115, 8.0),
            (11, "thruster_module", 125, 10.0),
            (12, "alloy_panel", 135, 8.0),
            (13, "circuit_board", 145, 8.0),
            (14, "habitat_unit", 155, 12.0),
            (15, "fuel_cell", 160, 8.0),
            (16, "solar_array", 175, 10.0),
            (17, "alloy_panel", 185, 8.0),
            (18, "thruster_module", 200, 10.0),
            (19, "circuit_board", 215, 8.0),
            (20, "habitat_unit", 235, 12.0),
        ]
    ]

    min_production_runs = 16
    min_assemblies = 10
    min_deliveries = 14
    min_energy_final = 15.0
    max_invalid_actions = 12


TASK_REGISTRY: Dict[str, Task] = {
    "easy": EasyTask(),
    "medium": MediumTask(),
    "hard": HardTask(),
}