src/tasks.py

from typing import Dict, Any, List
from .env import EmergencyResourceEnv


class Task:
    def __init__(self, name: str, difficulty: str, env: EmergencyResourceEnv):
        self.name = name
        self.difficulty = difficulty
        self.env = env
    
    def setup(self):
        """Setup task-specific environment configuration"""
        self.env.reset({
            'difficulty': self.difficulty,
            'task_name': self.name
        })


class EasyTask(Task):
    def __init__(self, env: EmergencyResourceEnv):
        super().__init__("Emergency Response Basic", "easy", env)
    
    def setup(self):
        self.env.reset({
            'difficulty': 'easy',
            'task_name': self.name,
            'expected_steps': 15
        })


class MediumTask(Task):
    def __init__(self, env: EmergencyResourceEnv):
        super().__init__("Emergency Response Advanced", "medium", env)
    
    def setup(self):
        self.env.reset({
            'difficulty': 'medium',
            'task_name': self.name,
            'expected_steps': 30
        })


class HardTask(Task):
    def __init__(self, env: EmergencyResourceEnv):
        super().__init__("Emergency Response Expert", "hard", env)
    
    def setup(self):
        self.env.reset({
            'difficulty': 'hard',
            'task_name': self.name,
            'expected_steps': 50
        })