visualizer.py

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from matplotlib.animation import FuncAnimation
from matplotlib.patches import Patch

from grid import Grid
from frontier import UNKNOWN, get_frontiers

COLORS = {
    "unknown":        "#E2E8F0",
    "true_obstacle":  "#CBD5E0",
    "free_known":     "#FFFFFF",
    "obstacle_known": "#2D3748",
    "scan":           "#BEE3F8",
    "visited_goal":   "#9AE6B4",
    "visited_front":  "#FBD38D",
    "frontier":       "#E9D8FD",
    "start":          "#48BB78",
    "goal":           "#FC8181",
    "rover":          "#E53E3E",
}


def _build_frame(ax, ground_truth, known_map, rover_pos, history,
                 mode_history, current_scan_cells, start, goal,
                 step_n, replans, current_mode):
    """
    Render a single animation frame onto ax.

    Args:
        ax: Matplotlib axes to draw on.
        ground_truth: Full Grid (used to reveal hidden obstacles).
        known_map: Incrementally updated Grid seen by the rover.
        rover_pos: Current (x, y) of the rover.
        history: List of (x, y) positions visited so far.
        mode_history: List of mode strings ('goal'/'frontier') per step.
        current_scan_cells: Cells in the current sensor FOV.
        start: (x, y) start cell.
        goal: (x, y) goal cell.
        step_n: Current step number (displayed in title).
        replans: Total replanning count (displayed in title).
        current_mode: Mode string for current step.
    """
    ax.clear()
    w, h = ground_truth.width, ground_truth.height
    color_grid = np.full((h, w, 3), 1.0)

    for y in range(h):
        for x in range(w):
            km = known_map.grid[y][x]
            gt = ground_truth.grid[y][x]
            if km == UNKNOWN:
                c = COLORS["true_obstacle"] if gt == 1 else COLORS["unknown"]
            elif km == 1:
                c = COLORS["obstacle_known"]
            else:
                c = COLORS["free_known"]
            color_grid[y][x] = mcolors.to_rgb(c)

    for (x, y) in get_frontiers(known_map):
        color_grid[y][x] = mcolors.to_rgb(COLORS["frontier"])

    for (x, y) in current_scan_cells:
        if ground_truth.grid[y][x] == 0:
            color_grid[y][x] = mcolors.to_rgb(COLORS["scan"])

    for i, (x, y) in enumerate(history):
        mode = mode_history[i] if i < len(mode_history) else 'goal'
        color_grid[y][x] = mcolors.to_rgb(
            COLORS["visited_goal"] if mode == 'goal' else COLORS["visited_front"]
        )

    sx, sy = start
    gx, gy = goal
    rx, ry = rover_pos
    color_grid[sy][sx] = mcolors.to_rgb(COLORS["start"])
    color_grid[gy][gx] = mcolors.to_rgb(COLORS["goal"])
    color_grid[ry][rx] = mcolors.to_rgb(COLORS["rover"])

    ax.imshow(color_grid, origin="upper")
    for x in range(w + 1):
        ax.axvline(x - 0.5, color="#A0AEC0", linewidth=0.3)
    for y in range(h + 1):
        ax.axhline(y - 0.5, color="#A0AEC0", linewidth=0.3)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.text(sx, sy, "S", ha="center", va="center", fontsize=7, fontweight="bold", color="white")
    ax.text(gx, gy, "G", ha="center", va="center", fontsize=7, fontweight="bold", color="white")

    mode_label = {"goal": "→ Goal", "frontier": "⊕ Exploring", "stuck": "✗ Stuck"}.get(current_mode, "")
    ax.set_title(
        f"Step {step_n}  |  Replans: {replans}  |  Mode: {mode_label}",
        color="#1A202C", fontsize=10, pad=6
    )


def animate(ground_truth, rover, interval: int = 120):
    """
    Replay the rover's journey as a matplotlib animation.

    Each frame corresponds to one navigation step. The known map is
    reconstructed incrementally from scan_history to match the rover's
    perspective at each point in time.

    Args:
        ground_truth: The full Grid (used only for visual hints).
        rover: A completed Rover instance with scan_history populated.
        interval: Milliseconds between animation frames.

    Returns:
        FuncAnimation object (or None if goal was unreachable).
    """
    fig, ax = plt.subplots(figsize=(8, 8))
    fig.patch.set_facecolor("#F7FAFC")
    fig.text(0.5, 0.97, "Autonomous Rover Navigator — Optimistic A* + Frontier Fallback",
             ha="center", va="top", fontsize=11, color="#1A202C", fontweight="bold")

    legend_patches = [
        Patch(color=COLORS["rover"],          label="Rover"),
        Patch(color=COLORS["visited_goal"],   label="Trail (to goal)"),
        Patch(color=COLORS["visited_front"],  label="Trail (exploring)"),
        Patch(color=COLORS["frontier"],       label="Frontier cells"),
        Patch(color=COLORS["scan"],           label="Sensor FOV"),
        Patch(color=COLORS["obstacle_known"], label="Obstacle (detected)"),
        Patch(color=COLORS["true_obstacle"],  label="Obstacle (hidden)"),
        Patch(color=COLORS["unknown"],        label="Unknown"),
    ]
    ax.legend(handles=legend_patches, loc="lower right", fontsize=7,
              framealpha=0.8, facecolor="white")

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

    if not rover.scan_history:
        _build_frame(ax, ground_truth, km, rover.start, [], [], [],
                     rover.start, rover.goal, 0, 0, 'stuck')
        fig.text(0.5, 0.02, "Goal is unreachable. Try a different --seed.",
                 ha="center", color="red", fontsize=11)
        plt.tight_layout(rect=[0, 0.04, 1, 0.95])
        plt.show()
        return None

    total_frames = len(rover.scan_history) + 1

    def update(frame):
        """Update callback for FuncAnimation."""
        if frame < len(rover.scan_history):
            pos, obs = rover.scan_history[frame]
            for (x, y), val in obs.items():
                km.grid[y][x] = val
            scan_cells = list(obs.keys())
            hist = rover.history[:frame]
            cur_pos = rover.history[frame] if frame < len(rover.history) else rover.position
            cur_mode = rover.mode_history[frame] if frame < len(rover.mode_history) else 'goal'
        else:
            scan_cells = []
            hist = rover.history
            cur_pos = rover.position
            cur_mode = rover.mode_history[-1] if rover.mode_history else 'goal'

        _build_frame(ax, ground_truth, km, cur_pos, hist,
                     rover.mode_history, scan_cells,
                     rover.start, rover.goal, frame,
                     rover.replans, cur_mode)

    anim = FuncAnimation(fig, update, frames=total_frames, interval=interval, repeat=False)
    plt.tight_layout(rect=[0, 0, 1, 0.95])
    plt.show()
    return anim