src/env.py

import numpy as np
from typing import Dict, List, Optional, Tuple, Any
from dataclasses import dataclass, asdict
from datetime import datetime
import random

from .models import Observation, Action, Reward, State, ServiceRequest


class EmergencyResourceEnv:
    """
    Emergency Resource Allocation Environment
    Grid-based city where agent allocates resources to service requests
    """
    
    def __init__(self, grid_size: int = 5, max_time_steps: int = 50):
        self.grid_size = grid_size
        self.max_time_steps = max_time_steps
        
        # Environment state
        self.agent_pos: Tuple[int, int] = (0, 0)
        self.requests: List[ServiceRequest] = []
        self.time_left: int = max_time_steps
        self.resources_left: int = 10
        self.current_step: int = 0
        self.total_reward: float = 0.0
        self.done: bool = False
        self.task_config: Dict[str, Any] = {}
        
    def reset(self, task_config: Optional[Dict[str, Any]] = None) -> Observation:
        """Reset environment to initial state"""
        self.agent_pos = (0, 0)
        self.time_left = self.max_time_steps
        self.resources_left = 10
        self.current_step = 0
        self.total_reward = 0.0
        self.done = False
        
        # Load task configuration
        self.task_config = task_config or {}
        
        # Generate service requests based on task
        self._generate_requests()
        
        return self._get_observation()
    
    def _generate_requests(self):
        """Generate service requests based on task difficulty"""
        task_difficulty = self.task_config.get('difficulty', 'medium')
        
        if task_difficulty == 'easy':
            num_requests = 3
            critical_ratio = 0.33  # 1 critical, 2 normal
        elif task_difficulty == 'medium':
            num_requests = 5
            critical_ratio = 0.4  # 2 critical, 3 normal
        else:  # hard
            num_requests = 8
            critical_ratio = 0.5  # 4 critical, 4 normal
        
        self.requests = []
        for i in range(num_requests):
            priority = 'critical' if random.random() < critical_ratio else 'normal'
            # Ensure critical requests aren't too clustered
            if priority == 'critical' and len([r for r in self.requests if r.priority == 'critical']) >= num_requests * critical_ratio:
                priority = 'normal'
                
            # Generate position not overlapping with existing requests
            while True:
                pos = (random.randint(0, self.grid_size-1), 
                       random.randint(0, self.grid_size-1))
                if not any(r.position == pos for r in self.requests):
                    break
            
            self.requests.append(ServiceRequest(
                id=i,
                position=pos,
                priority=priority,
                allocated=False,
                created_at=self.current_step
            ))
    
    def step(self, action: Action) -> Tuple[Observation, Reward, bool, Dict]:
        """Execute action and return new state"""
        if self.done:
            return self._get_observation(), Reward(score=0.0), True, {"error": "Episode already done"}
        
        # Execute action
        action_type = action.action_type
        error = None
        
        try:
            if action_type == "move":
                self._move(action.direction)
            elif action_type == "allocate":
                reward_change = self._allocate_resources()
            else:
                error = f"Unknown action type: {action_type}"
                reward_change = -1.0
        except Exception as e:
            error = str(e)
            reward_change = -1.0
        
        # Apply step penalty
        step_penalty = -0.5
        self.total_reward += step_penalty
        self.current_step += 1
        self.time_left -= 1
        
        # Check episode termination
        all_requests_served = all(r.allocated for r in self.requests)
        time_out = self.time_left <= 0
        resources_depleted = self.resources_left <= 0
        
        if all_requests_served:
            self.total_reward += 30  # Completion bonus
            self.done = True
        elif time_out:
            self.total_reward -= 20  # Time out penalty
            self.done = True
        elif resources_depleted:
            self.done = True
        elif self.current_step >= self.max_time_steps:
            self.done = True
        
        # Get final reward
        reward = Reward(score=max(0.0, min(1.0, (self.total_reward + 50) / 100)))
        
        return self._get_observation(), reward, self.done, {"error": error}
    
    def _move(self, direction: str):
        """Move agent in grid"""
        x, y = self.agent_pos
        
        if direction == "up" and x > 0:
            self.agent_pos = (x - 1, y)
        elif direction == "down" and x < self.grid_size - 1:
            self.agent_pos = (x + 1, y)
        elif direction == "left" and y > 0:
            self.agent_pos = (x, y - 1)
        elif direction == "right" and y < self.grid_size - 1:
            self.agent_pos = (x, y + 1)
    
    def _allocate_resources(self) -> float:
        """Allocate resources at current position"""
        reward_change = 0.0
        
        # Check if any request at current position
        for request in self.requests:
            if not request.allocated and request.position == self.agent_pos:
                if self.resources_left > 0:
                    request.allocated = True
                    self.resources_left -= 1
                    
                    if request.priority == "critical":
                        reward_change = 20.0
                    else:
                        reward_change = 10.0
                    
                    self.total_reward += reward_change
                break
        
        return reward_change
    
    def _get_observation(self) -> Observation:
        """Get current observation"""
        # Create grid representation
        grid = np.zeros((self.grid_size, self.grid_size, 3))  # 3 channels: agent, critical, normal
        
        # Agent position
        grid[self.agent_pos[0], self.agent_pos[1], 0] = 1
        
        # Requests
        for req in self.requests:
            if not req.allocated:
                channel = 1 if req.priority == "critical" else 2
                grid[req.position[0], req.position[1], channel] = 1
        
        return Observation(
            agent_position=self.agent_pos,
            grid=grid.tolist(),
            requests=[r for r in self.requests if not r.allocated],
            time_left=self.time_left,
            resources_left=self.resources_left,
            step=self.current_step,
            task_config=self.task_config
        )
    
    def state(self) -> State:
        """Get current state"""
        return State(
            agent_position=self.agent_pos,
            requests=[asdict(r) for r in self.requests],
            time_left=self.time_left,
            resources_left=self.resources_left,
            current_step=self.current_step,
            done=self.done,
            total_reward=self.total_reward
        )