visualize.py

import sys
sys.path.insert(0, ".")

import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.animation as animation
import numpy as np

from server.trafficops_environment import TrafficOpsEnvironment
from models import TrafficOpsAction


def run_episode(task="incident_corridor", seed=42):
    env = TrafficOpsEnvironment()
    obs = env.reset(seed=seed, task=task)
    frames = [obs]

    # Simple scripted agent
    while not obs.done:
        action = TrafficOpsAction(op="noop")

        # Reroute around incidents
        if obs.incidents:
            for inc in obs.incidents:
                if inc.active and not any(p.op == "reroute" for p in obs.active_plans):
                    if inc.road_id == "R_h_1_1":
                        action = TrafficOpsAction(
                            op="reroute", targets=[inc.road_id],
                            params={"blocked_road": inc.road_id, "detour": ["R_v_1_1", "R_h_2_1", "R_v_2_2"], "duration_ticks": 200},
                            reason=f"reroute around {inc.id}")
                        break

        # Preempt for emergencies
        if action.op == "noop" and obs.emergencies:
            for em in obs.emergencies:
                if not em.cleared and em.remaining_route:
                    targets = []
                    direction = None
                    for rid in em.remaining_route[:3]:
                        for iv in obs.intersections:
                            if any(rid == r for r in [iv.id]):
                                continue
                        for road in obs.roads:
                            if road.id == rid:
                                direction = road.approach_direction
                                targets.append(road.to_node)
                    targets = [t for t in targets if t.startswith("I_")][:3]
                    if targets and direction:
                        action = TrafficOpsAction(
                            op="preempt", targets=targets,
                            params={"direction": direction, "duration_ticks": 15},
                            reason=f"clear path for {em.id}")
                        break

        obs = env.step(action)
        frames.append(obs)

    return frames


def draw_frame(ax, obs, frame_num):
    ax.clear()
    ax.set_xlim(-1.5, 5.5)
    ax.set_ylim(-1.5, 5.5)
    ax.set_aspect("equal")
    ax.set_facecolor("#1a1a2e")

    # Title
    score_str = f"  Score: {obs.final_score:.3f}" if obs.final_score else ""
    ax.set_title(
        f"TrafficOps — {obs.task}  |  Tick {obs.tick}/{obs.horizon}  |  "
        f"Budget: {obs.interventions_used}/{obs.interventions_budget}{score_str}",
        color="white", fontsize=11, fontweight="bold", pad=10
    )

    # Draw roads as gray lines
    for road in obs.roads:
        from_pos = None
        to_pos = None
        for iv in obs.intersections:
            if iv.id == road.from_node:
                from_pos = iv.position
            if iv.id == road.to_node:
                to_pos = iv.position
        if from_pos and to_pos:
            color = "#dc2626" if road.blocked else "#4a4a6a"
            lw = 3 if road.blocked else 1.5
            ax.plot([from_pos[0], to_pos[0]], [from_pos[1], to_pos[1]],
                    color=color, linewidth=lw, zorder=1)
            # Show occupancy as dots along road
            if road.occupancy > 0:
                mid_x = (from_pos[0] + to_pos[0]) / 2
                mid_y = (from_pos[1] + to_pos[1]) / 2
                ax.annotate(str(road.occupancy), (mid_x, mid_y),
                           color="#94a3b8", fontsize=6, ha="center", va="center")

    # Draw source/sink arrows
    for road in obs.roads:
        if road.from_node.startswith("SRC"):
            to_pos = None
            for iv in obs.intersections:
                if iv.id == road.to_node:
                    to_pos = iv.position
            if to_pos:
                if "W" in road.from_node:
                    ax.annotate("→", (to_pos[0] - 1.2, to_pos[1]), color="#64748b", fontsize=10, ha="center", va="center")
                elif "S" in road.from_node:
                    ax.annotate("↑", (to_pos[0], to_pos[1] - 1.2), color="#64748b", fontsize=10, ha="center", va="center")

    # Draw intersections
    for iv in obs.intersections:
        x, y = iv.position
        is_ns = "N" in iv.current_phase

        # Signal phase colors
        if iv.preempt_direction:
            color = "#fbbf24"  # yellow for preempt
            edge = "#f59e0b"
        elif is_ns:
            color = "#22c55e"  # green for NS
            edge = "#16a34a"
        else:
            color = "#3b82f6"  # blue for EW
            edge = "#2563eb"

        circle = plt.Circle((x, y), 0.35, facecolor=color, edgecolor=edge, linewidth=2, zorder=5)
        ax.add_patch(circle)

        # Queue bars
        q_s = iv.queues.get("S", 0)
        q_w = iv.queues.get("W", 0)
        q_n = iv.queues.get("N", 0)
        q_e = iv.queues.get("E", 0)

        bar_scale = 0.08
        max_bar = 0.8
        if q_s > 0:
            h = min(q_s * bar_scale, max_bar)
            ax.add_patch(patches.Rectangle((x - 0.1, y - 0.4 - h), 0.2, h,
                         facecolor="#ef4444", alpha=0.8, zorder=4))
        if q_w > 0:
            w = min(q_w * bar_scale, max_bar)
            ax.add_patch(patches.Rectangle((x - 0.4 - w, y - 0.1), w, 0.2,
                         facecolor="#f97316", alpha=0.8, zorder=4))
        if q_n > 0:
            h = min(q_n * bar_scale, max_bar)
            ax.add_patch(patches.Rectangle((x - 0.1, y + 0.4), 0.2, h,
                         facecolor="#ef4444", alpha=0.6, zorder=4))
        if q_e > 0:
            w = min(q_e * bar_scale, max_bar)
            ax.add_patch(patches.Rectangle((x + 0.4, y - 0.1), w, 0.2,
                         facecolor="#f97316", alpha=0.6, zorder=4))

        # Intersection label
        label = iv.id.replace("I_", "")
        ax.text(x, y, label, color="white", fontsize=7, ha="center", va="center",
                fontweight="bold", zorder=6)

        # Preempt indicator
        if iv.preempt_direction:
            ax.text(x, y + 0.55, f"⚡{iv.preempt_direction}", color="#fbbf24",
                    fontsize=7, ha="center", va="center", zorder=6)

    # Draw emergencies
    for em in obs.emergencies:
        if em.cleared:
            continue
        for iv in obs.intersections:
            for road in obs.roads:
                if road.id == em.current_road and road.to_node == iv.id:
                    x, y = iv.position
                    emoji = {"ambulance": "🚑", "fire": "🚒", "police": "🚓"}.get(em.type, "🚨")
                    ax.text(x + 0.5, y + 0.5, emoji, fontsize=12, ha="center", va="center", zorder=10)

    # Draw incidents
    for inc in obs.incidents:
        if not inc.active:
            continue
        ax.text(0.02, 0.02 + obs.incidents.index(inc) * 0.04,
                f"🚧 {inc.id}: {inc.kind} blocks {inc.road_id}",
                transform=ax.transAxes, color="#ef4444", fontsize=8, va="bottom")

    # Metrics bar at bottom
    m = obs.metrics
    if m:
        metrics_text = (
            f"Cleared: {m.cleared_civilian}/{m.spawned_civilian} civ  "
            f"{m.cleared_emergency}/{m.spawned_emergency} em  |  "
            f"Wait: avg={m.mean_wait_ticks:.0f} max={m.max_wait_ticks}  |  "
            f"Wasted: {m.wasted_green_ticks}  |  Gridlocks: {m.gridlock_events}"
        )
        ax.text(0.5, -0.08, metrics_text, transform=ax.transAxes,
                color="#94a3b8", fontsize=8, ha="center", va="top")

    # Legend
    legend_items = [
        ("●", "#22c55e", "N-S Green"),
        ("●", "#3b82f6", "E-W Green"),
        ("●", "#fbbf24", "Preempt"),
        ("■", "#ef4444", "Queue S"),
        ("■", "#f97316", "Queue W"),
        ("—", "#dc2626", "Blocked"),
    ]
    for i, (sym, col, label) in enumerate(legend_items):
        ax.text(4.5 + (i // 3) * 1.5, -1.0 - (i % 3) * 0.3, f"{sym} {label}",
                color=col, fontsize=7, va="center")

    ax.axis("off")


def main():
    task = sys.argv[1] if len(sys.argv) > 1 else "incident_corridor"
    seed = int(sys.argv[2]) if len(sys.argv) > 2 else 42

    print(f"Running {task} (seed={seed})...")
    frames = run_episode(task, seed)
    print(f"Got {len(frames)} frames, final_score={frames[-1].final_score}")

    fig, ax = plt.subplots(1, 1, figsize=(10, 10))
    fig.patch.set_facecolor("#0f0f23")

    def update(frame_num):
        if frame_num < len(frames):
            draw_frame(ax, frames[frame_num], frame_num)
        return []

    ani = animation.FuncAnimation(fig, update, frames=len(frames), interval=300, repeat=True)

    # Save as GIF
    print("Saving trafficops_demo.gif...")
    ani.save("trafficops_demo.gif", writer="pillow", fps=4, dpi=100)
    print("Done! Open trafficops_demo.gif")

    plt.show()


if __name__ == "__main__":
    main()