rover.py

from grid import Grid
from sensor import LidarSensor
from astar import astar
from frontier import nearest_frontier, UNKNOWN
import numpy as np


class OptimisticGrid:
    """Wrapper that treats UNKNOWN cells as free for optimistic planning."""

    def __init__(self, known_map):
        self.km = known_map
        self.width = known_map.width
        self.height = known_map.height
        self.grid = known_map.grid  # for visualizer compatibility

    def get_neighbors(self, x, y):
        directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]
        result = []
        for dx, dy in directions:
            nx, ny = x+dx, y+dy
            if 0 <= nx < self.width and 0 <= ny < self.height:
                if self.km.grid[ny][nx] != 1:  # free or unknown = passable
                    result.append((nx, ny))
        return result


class Rover:
    """
    Autonomous rover using Optimistic A*:
    - Unknown cells assumed free (optimistic)
    - Replans only when obstacle found on current path
    - Falls back to frontier exploration if goal appears unreachable
    """

    def __init__(self, grid: Grid, start: tuple, goal: tuple, sensor_radius: int = 4):
        self.ground_truth = grid
        self.start = start
        self.goal = goal
        self.position = start
        self.sensor = LidarSensor(radius=sensor_radius)

        self.known_map = Grid(grid.width, grid.height, obstacle_ratio=0.0)
        self.known_map.grid = np.full((grid.height, grid.width), UNKNOWN, dtype=int)

        self.path = []
        self.history = []
        self.scan_history = []
        self.replans = 0
        self.mode_history = []
        self.current_target = None

    def _check_reachability(self):
        path, _ = astar(self.ground_truth, self.position, self.goal)
        return path is not None

    def _obstacle_on_path(self):
        """Check if any known obstacle appeared on the upcoming path."""
        for (x, y) in self.path[1:self.sensor.radius + 2]:
            if self.known_map.grid[y][x] == 1:
                return True
        return False

    def _plan_optimistic(self, target):
        """Plan to target treating unknowns as free."""
        og = OptimisticGrid(self.known_map)
        path, _ = astar(og, self.position, target)
        return path

    def _replan(self):
        # Primary: optimistic path to goal
        path = self._plan_optimistic(self.goal)
        if path:
            self.current_target = self.goal
            self.mode_history.append('goal')
            return path

        # Fallback: frontier exploration
        f = nearest_frontier(self.known_map, self.position)
        if f:
            path = self._plan_optimistic(f)
            if path:
                self.current_target = f
                self.mode_history.append('frontier')
                return path

        self.mode_history.append('stuck')
        return None

    def step(self):
        if self.position == self.goal:
            return False

        obs = self.sensor.scan(self.ground_truth, *self.position)
        self.scan_history.append((self.position, obs))

        map_changed = False
        for (x, y), val in obs.items():
            if self.known_map.grid[y][x] != val:
                self.known_map.grid[y][x] = val
                map_changed = True

        needs_replan = (
            not self.path
            or self.position == self.current_target
            or (map_changed and self._obstacle_on_path())
        )

        if needs_replan:
            self.path = self._replan()
            self.replans += 1
        else:
            self.mode_history.append('goal' if self.current_target == self.goal else 'frontier')

        if not self.path or len(self.path) < 2:
            return False

        self.history.append(self.position)
        self.position = self.path[1]
        self.path = self.path[1:]
        return True

    def run(self, max_steps=5000):
        if not self._check_reachability():
            return 'unreachable'
        for _ in range(max_steps):
            if not self.step():
                break
        return 'success' if self.position == self.goal else 'stuck'