server/grid_world_environment.py

"""
server/grid_world_environment.py
Core Drone Delivery RL Environment — refactored to use core/ modules.
"""
from __future__ import annotations

import random
from typing import Any, Dict, List, Optional, Tuple

import torch

try:
    from openenv.core.env_server import Environment
except ImportError:
    class Environment:
        def reset(self): raise NotImplementedError
        def step(self, action): raise NotImplementedError
        @property
        def state(self): raise NotImplementedError

from drone_env.models import DroneAction, DroneObservation, DroneState
from drone_env.core.grid_generator import generate_city_map, EMOJI, LEGEND
from drone_env.core.tasks import TASK_CONFIG
from drone_env.core.drone import compute_next_pos, drain_battery
from drone_env.core.obstacles import check_move
from drone_env.core.state_manager import new_episode_state
try:
    from drone_env.graders import GRADERS
except ImportError:
    try:
        from graders import GRADERS
    except ImportError:
        # Final fallback
        GRADERS = {}
from drone_env.rl.trainer import record_episode


class DroneDeliveryEnvironment(Environment):

    def __init__(self):
        super().__init__()
        self._state = DroneState()
        self._cfg = TASK_CONFIG["graders:grade_easy"]
        self._grid: List[List[str]] = []
        self._deliveries: List[Tuple[int, int]] = []
        self._delivered: List[bool] = []
        self._drone_x = 0
        self._drone_y = 0
        self._battery = 0
        self._rng = torch.Generator()
        self._step_records: List[Dict] = []
        self._start_pos: Tuple[int, int] = (0, 0)

    def reset(self, action: Optional[DroneAction] = None) -> DroneObservation:
        # Potentially persist old episode if not already done
        if self._step_records and self._state.episode_id and not self._state.done:
             self._persist_episode()

        # Map short task IDs (from openenv.yaml) to full grader keys
        _TASK_ID_MAP = {
            "easy_delivery":   "graders:grade_easy",
            "medium_delivery": "graders:grade_medium",
            "hard_delivery":   "graders:grade_hard",
            # Legacy redirects
            "drone_env.graders.easy:grade_easy":   "graders:grade_easy",
            "drone_env.graders.medium:grade_medium": "graders:grade_medium",
            "drone_env.graders.hard:grade_hard":   "graders:grade_hard",
        }
        task = "graders:grade_easy"
        if action and action.task_name:
            name = action.task_name
            if name in _TASK_ID_MAP:
                task = _TASK_ID_MAP[name]
            elif name in TASK_CONFIG:
                task = name
            else:
                # handle partial matches like "easy"
                for k in TASK_CONFIG:
                    if k.startswith(name):
                        task = k; break

        self._cfg = TASK_CONFIG[task]
        self._rng.manual_seed(random.randint(0, 2**31))

        self._grid, self._deliveries, self._start_pos = generate_city_map(self._cfg, self._rng)
        self._delivered = [False] * len(self._deliveries)
        self._drone_x, self._drone_y = self._start_pos
        self._battery = self._cfg["battery_max"]
        self._step_records = []

        self._state = new_episode_state(
            task=task,
            deliveries=self._deliveries,
            start_x=self._drone_x,
            start_y=self._drone_y,
            battery_max=self._cfg["battery_max"],
        )
        return self._build_obs(0.0, "🚁 Mission started! Deliver all packages.")

    def step(self, action: DroneAction) -> DroneObservation:
        if self._state.done:
            return self._build_obs(0.0, "Episode ended. Call reset().")

        cfg = self._cfg
        direction = (action.direction or "WAIT").upper()
        nx, ny = compute_next_pos(self._drone_x, self._drone_y, direction)
        outcome, cell = check_move(self._grid, nx, ny, cfg["width"], cfg["height"])

        reward = cfg["r_step"]
        msg = ""

        # Reward shaping (Distance bonus)
        _, dist_before = self._next_target()
        
        if outcome == "wall":
            nx, ny = self._drone_x, self._drone_y
            reward += cfg["r_wall"]
            msg = f"Hit {direction} boundary! 🚫"
        elif outcome == "blocked":
            nx, ny = self._drone_x, self._drone_y
            reward += cfg["r_blocked"]
            msg = f"Path is {direction}ly blocked!"
        elif outcome == "building":
            self._drone_x, self._drone_y = nx, ny
            reward = cfg.get("r_building", -0.1)
            msg = f"Flying above a building! {EMOJI['building']}"
        elif outcome == "tree":
            self._drone_x, self._drone_y = nx, ny
            reward = cfg.get("r_tree", -0.1)
            msg = f"Flying over a tree! {EMOJI['tree']}"
        elif outcome == "obstacle":
            self._drone_x, self._drone_y = nx, ny
            reward += cfg["r_obstacle"]
            msg = f"Hit obstacle! {EMOJI['obstacle']}"
        else:
            self._drone_x, self._drone_y = nx, ny
            if direction == "WAIT":
                reward = cfg.get("r_wait", -0.1)
                msg = "Drone is idling (WAIT)... 🔋"
            elif not msg:
                # Get cell type at current position
                current_cell = self._grid[self._drone_y][self._drone_x]
                if current_cell == "road":
                    msg = f"On road {EMOJI['road']}"
                else:
                    msg = f"Flying over {current_cell} {EMOJI.get(current_cell, '')}"

        # Move finalized, now check new distance
        _, dist_after = self._next_target()
        if dist_after is not None and dist_before is not None and outcome not in ["tree", "building", "ok"]:
            # Reward for moving closer, small penalty for moving away
            delta = dist_before - dist_after
            reward += delta * 0.05

        # Delivery check
        for i, (tx, ty) in enumerate(self._deliveries):
            if not self._delivered[i] and self._drone_x == tx and self._drone_y == ty:
                self._delivered[i] = True
                reward += cfg["r_delivery"]
                self._state.deliveries_done += 1
                msg = f"✅ Delivery {self._state.deliveries_done}/{len(self._deliveries)} done!"
                break

        # Battery drain
        self._battery = drain_battery(self._battery, cfg["battery_cost"])
        self._state.step_count += 1
        self._state.reward_total += reward
        bat_norm = max(0.0, float(self._battery) / cfg["battery_max"])
        self._state.battery = bat_norm
        self._state.drone_x = self._drone_x
        self._state.drone_y = self._drone_y

        self._step_records.append(dict(
            step=self._state.step_count,
            x=self._drone_x, y=self._drone_y,
            action=direction,
            reward=float(round(reward, 5)),
            battery=float(round(bat_norm, 4)),
            message=msg,
        ))
        self._state.path_history = self._step_records

        # Done conditions
        done = False
        if all(self._delivered):
            done = True; msg = "🎉 All deliveries complete!"
        elif self._battery <= 0:
            reward += cfg["r_battery_dead"]
            self._state.reward_total += cfg["r_battery_dead"]
            done = True; msg = "🔋 Battery dead!"
        elif self._state.step_count >= cfg["max_steps"]:
            done = True; msg = "⏰ Max steps reached."

        self._state.done = done
        if done:
            self._persist_episode()

        return self._build_obs(reward, msg)

    @property
    def state(self) -> DroneState:
        return self._state

    @property
    def graders(self) -> Dict:
        """Expose graders for the environment."""
        return GRADERS

    def _persist_episode(self):
        try:
            cfg = self._cfg
            record_episode(
                task=self._state.task_name,
                steps=list(self._step_records),
                grid_meta=dict(
                    width=cfg["width"], height=cfg["height"]
                ),
                delivery_positions=[[d[0], d[1]] for d in self._deliveries],
                deliveries_done=self._state.deliveries_done,
                total_reward=float(round(self._state.reward_total, 4)),
            )
        except Exception as e:
            print(f"[DroneEnv] Episode record error: {e}")

    def _build_obs(self, reward: float, message: str) -> DroneObservation:
        cfg = self._cfg
        target, dist = self._next_target()
        return DroneObservation(
            grid=self._render_grid(),
            grid_width=int(cfg["width"]), grid_height=int(cfg["height"]),
            drone_x=self._drone_x, drone_y=self._drone_y,
            battery=float(max(0.0, float(self._battery) / cfg["battery_max"])),
            battery_steps_remaining=int(self._battery),
            deliveries_total=int(len(self._deliveries)),
            deliveries_done=int(self._state.deliveries_done),
            current_target=target,
            distance_to_target=float(dist) if dist is not None else None,
            step_count=int(self._state.step_count),
            max_steps=int(cfg["max_steps"]),
            reward_last=float(max(0.01, min(0.99, round(reward, 4)))),
            reward_total=float(max(0.01, min(0.99, round(self._state.reward_total, 4)))),
            score=float(max(0.01, min(0.99, round(GRADERS[self._state.task_name](self._state), 4)))),
            done=bool(self._state.done),
            message=str(message),
            legend=dict(LEGEND),
            cell_types=list(self._grid),
        )

    def _render_grid(self) -> List[str]:
        rows = []
        for y, row in enumerate(self._grid):
            line = ""
            for x, cell in enumerate(row):
                if x == self._drone_x and y == self._drone_y:
                    line += EMOJI["drone"]
                    continue
                
                # Check for delivery targets
                is_delivery = False
                for i, (tx, ty) in enumerate(self._deliveries):
                    if x == tx and y == ty:
                        line += EMOJI["done_del"] if self._delivered[i] else EMOJI["delivery"]
                        is_delivery = True
                        break
                
                if not is_delivery:
                    # Map the internal symbol to emoji
                    # Now it looks up from EMOJI dict using the symbolic string in 'cell'
                    line += EMOJI.get(cell, cell)
            rows.append(line)
        return rows

    def _next_target(self):
        for i, (tx, ty) in enumerate(self._deliveries):
            if not self._delivered[i]:
                dist = float(abs(self._drone_x - tx) + abs(self._drone_y - ty))
                return (int(tx), int(ty)), float(dist)
        return None, None