fsd_model/visualization.py

"""
Visualization utilities for FSD model outputs.
Generates visual representations of:
- Sensor placement on vehicle
- BEV perception outputs
- Planned trajectory
- Control commands
"""

import torch
import numpy as np
from typing import Dict, Optional, Tuple, List
import math


def visualize_sensor_layout_ascii(vehicle_config) -> str:
    """Generate ASCII art of sensor placement on vehicle."""
    sc = vehicle_config.sensor_config
    half_l = vehicle_config.length / 2
    half_w = vehicle_config.width / 2
    
    # Create grid
    grid_h, grid_w = 30, 50
    grid = [[' '] * grid_w for _ in range(grid_h)]
    
    # Scale factors
    scale_x = (grid_w - 10) / (vehicle_config.length + 2)
    scale_y = (grid_h - 6) / (vehicle_config.width + 2)
    
    cx, cy = grid_w // 2, grid_h // 2
    
    # Draw vehicle outline
    vw = int(vehicle_config.length * scale_x / 2)
    vh = int(vehicle_config.width * scale_y / 2)
    
    for x in range(cx - vw, cx + vw + 1):
        if 0 <= x < grid_w:
            if 0 <= cy - vh < grid_h:
                grid[cy - vh][x] = '─'
            if 0 <= cy + vh < grid_h:
                grid[cy + vh][x] = '─'
    
    for y in range(cy - vh, cy + vh + 1):
        if 0 <= y < grid_h:
            if 0 <= cx - vw < grid_w:
                grid[y][cx - vw] = '│'
            if 0 <= cx + vw < grid_w:
                grid[y][cx + vw] = '│'
    
    # Corners
    for (dy, dx), ch in [
        ((-vh, -vw), '┌'), ((-vh, vw), '┐'),
        ((vh, -vw), '└'), ((vh, vw), '┘')
    ]:
        gy, gx = cy + dy, cx + dx
        if 0 <= gy < grid_h and 0 <= gx < grid_w:
            grid[gy][gx] = ch
    
    # Direction arrow
    if 0 <= cy < grid_h and cx + vw + 1 < grid_w:
        grid[cy][cx + vw + 1] = '►'
    grid[cy][cx] = '+'  # center
    
    # Place cameras
    for i, cam in enumerate(sc.cameras):
        gx = cx + int(cam.placement.x * scale_x)
        gy = cy - int(cam.placement.y * scale_y)
        if 0 <= gy < grid_h and 0 <= gx < grid_w:
            grid[gy][gx] = 'C'
    
    # Place ultrasonics
    for i, us in enumerate(sc.ultrasonics):
        gx = cx + int(us.placement.x * scale_x)
        gy = cy - int(us.placement.y * scale_y)
        if 0 <= gy < grid_h and 0 <= gx < grid_w:
            if grid[gy][gx] == ' ' or grid[gy][gx] in '─│':
                grid[gy][gx] = 'U'
    
    # Add labels
    result = "Vehicle Sensor Layout (Top View)\n"
    result += "C = Camera, U = Ultrasonic, + = Center, ► = Forward\n"
    result += "=" * grid_w + "\n"
    for row in grid:
        result += ''.join(row) + "\n"
    result += "=" * grid_w + "\n"
    
    return result


def format_model_output(output: Dict[str, torch.Tensor]) -> str:
    """Format model output as human-readable text."""
    lines = []
    lines.append("╔══════════════════════════════════════════╗")
    lines.append("║       FSD Model Output Summary           ║")
    lines.append("╠══════════════════════════════════════════╣")
    
    # Control outputs
    if "control/steering_deg" in output:
        steer = output["control/steering_deg"].mean().item()
        throttle = output["control/throttle"].mean().item()
        brake = output["control/brake"].mean().item()
        lines.append(f"║ Steering:  {steer:+.2f}°")
        lines.append(f"║ Throttle:  {throttle:.3f}")
        lines.append(f"║ Brake:     {brake:.3f}")
    
    # Safety
    if "planning/collision_risk" in output:
        risk = output["planning/collision_risk"].mean().item()
        emergency = output["planning/emergency_brake"].mean().item()
        lines.append(f"║ Collision Risk: {risk:.3f}")
        lines.append(f"║ Emergency Brake: {emergency:.3f}")
    
    # Behavior
    if "planning/behavior_logits" in output:
        behaviors = [
            "keep_lane", "turn_left", "turn_right", "change_left",
            "change_right", "stop", "yield", "park", "reverse", "emergency"
        ]
        probs = torch.softmax(output["planning/behavior_logits"], dim=-1)
        top_prob, top_idx = probs.mean(0).topk(3)
        lines.append("║ Top behaviors:")
        for p, i in zip(top_prob, top_idx):
            lines.append(f"║   {behaviors[i.item()]}: {p.item():.3f}")
    
    # Waypoints
    if "planning/safe_waypoints" in output:
        wp = output["planning/safe_waypoints"]
        lines.append(f"║ Planned waypoints: {wp.shape[1]}")
        lines.append(f"║   First: ({wp[0,0,0]:.2f}, {wp[0,0,1]:.2f})")
        lines.append(f"║   Last:  ({wp[0,-1,0]:.2f}, {wp[0,-1,1]:.2f})")
    
    # Controller weights
    if "control/controller_weights" in output:
        w = output["control/controller_weights"].mean(0)
        lines.append(f"║ Controller mix: Neural={w[0]:.2f} Stanley={w[1]:.2f} PID={w[2]:.2f}")
    
    lines.append("╚══════════════════════════════════════════╝")
    
    return "\n".join(lines)


def format_parameter_count(counts: Dict[str, int]) -> str:
    """Format parameter counts as a nice table."""
    lines = []
    lines.append("┌─────────────────┬──────────────┬───────────┐")
    lines.append("│ Module          │ Parameters   │ Size (MB) │")
    lines.append("├─────────────────┼──────────────┼───────────┤")
    
    for name, count in counts.items():
        if name in ["total", "total_trainable"]:
            continue
        size_mb = count * 4 / (1024 * 1024)  # float32
        lines.append(f"│ {name:<15} │ {count:>12,} │ {size_mb:>8.2f}  │")
    
    lines.append("├─────────────────┼──────────────┼───────────┤")
    total = counts.get("total", 0)
    trainable = counts.get("total_trainable", 0)
    total_mb = total * 4 / (1024 * 1024)
    lines.append(f"│ {'TOTAL':<15} │ {total:>12,} │ {total_mb:>8.2f}  │")
    lines.append(f"│ {'Trainable':<15} │ {trainable:>12,} │          │")
    lines.append("└─────────────────┴──────────────┴───────────┘")
    
    return "\n".join(lines)