test_trajectory_logic.py

#!/usr/bin/env python3
"""
Test script to validate the trajectory visualization logic without running matplotlib.
This checks if the data loading and processing functions work correctly.
"""

import json
import os
import re

def test_get_parcel_color_families():
    """Test parcel color family function"""
    color_families = {
        '0': [0.2, 0.4, 0.8],    # Blue family
        '1': [0.8, 0.2, 0.2],    # Red family  
        '2': [0.2, 0.8, 0.2],    # Green family
        '3': [0.8, 0.6, 0.2],    # Orange family
        '4': [0.6, 0.2, 0.8],    # Purple family
        '5': [0.8, 0.8, 0.2],    # Yellow family
        '6': [0.2, 0.8, 0.8],    # Cyan family
        '7': [0.8, 0.2, 0.6],    # Magenta family
        '8': [0.4, 0.8, 0.4],    # Light Green family
        '9': [0.8, 0.4, 0.2],    # Brown family
    }
    print(f"✓ Color families defined for {len(color_families)} parcels")
    return color_families

def test_speed_to_color_intensity():
    """Test speed to color intensity conversion"""
    base_color = [0.2, 0.4, 0.8]  # Blue
    min_speed = 0.0
    max_speed = 0.05
    
    test_speeds = [0.0, 0.01, 0.025, 0.04, 0.05]
    print("\nSpeed to Color Intensity Test:")
    for speed in test_speeds:
        if max_speed <= min_speed:
            normalized_speed = 0.5
        else:
            normalized_speed = max(0.0, min(1.0, (speed - min_speed) / (max_speed - min_speed)))
        
        intensity = 0.3 + (0.7 * normalized_speed)
        color = [c * intensity for c in base_color]
        print(f"  Speed {speed:.3f} → Intensity {intensity:.3f} → Color {color}")
    
    print("✓ Speed to color intensity conversion working")

def test_filename_parsing():
    """Test filename parsing logic"""
    test_files = [
        "robot0_parcel0_approach_robot_trajectory.json",
        "robot1_parcel2_approach_robot_trajectory.json",
        "robot4_parcel5_approach_robot_trajectory.json",
        "invalid_filename.json"
    ]
    
    print("\nFilename Parsing Test:")
    for filename in test_files:
        match = re.match(r'robot(\d+)_parcel(\d+)_approach_robot_trajectory\.json', filename)
        if match:
            robot_id = match.group(1)
            parcel_id = match.group(2)
            print(f"  {filename} → Robot {robot_id}, Parcel {parcel_id} ✓")
        else:
            print(f"  {filename} → No match ✗")
    
    print("✓ Filename parsing logic working")

def test_trajectory_data_structure():
    """Test expected trajectory data structure"""
    # Simulate the expected data structure from the JSON files
    sample_data = {
        "description": "ApproachPlan robot real trajectory data [timestamp, x, y, theta, v, omega]",
        "data": [
            [1757425031.317215, 0.08625, 0.94873, -0.82918, 0.00123, -0.04595],
            [1757425031.417507, 0.08623, 0.94889, -0.81470, 0.01843, -0.09575],
            [1757425031.516950, 0.08631, 0.94873, -0.82933, 0.02041, -0.08702]
        ]
    }
    
    print("\nTrajectory Data Processing Test:")
    trajectory_points = []
    
    if 'data' in sample_data and isinstance(sample_data['data'], list):
        for point in sample_data['data']:
            if len(point) >= 6:
                timestamp, x, y, theta, v, omega = point[:6]
                trajectory_points.append({
                    'timestamp': timestamp,
                    'x': x,
                    'y': y, 
                    'theta': theta,
                    'velocity': abs(v),
                    'omega': omega
                })
    
    print(f"  Processed {len(trajectory_points)} trajectory points")
    for i, point in enumerate(trajectory_points):
        print(f"    Point {i}: x={point['x']:.5f}, y={point['y']:.5f}, v={point['velocity']:.5f}")
    
    print("✓ Trajectory data processing working")

def verify_data_files():
    """Check if actual trajectory files exist"""
    data_dir = os.path.join(os.path.dirname(__file__), "result", "romi", "success0", "experi", "control_data")
    pattern = os.path.join(data_dir, "robot*_parcel*_approach_robot_trajectory.json")
    
    print(f"\nData Files Verification:")
    print(f"Looking in: {data_dir}")
    
    if os.path.exists(data_dir):
        print("✓ Data directory exists")
        
        # Count trajectory files
        import glob
        trajectory_files = glob.glob(pattern)
        print(f"✓ Found {len(trajectory_files)} trajectory files")
        
        if trajectory_files:
            # Test loading one file
            test_file = trajectory_files[0]
            print(f"Testing file: {os.path.basename(test_file)}")
            
            try:
                with open(test_file, 'r') as f:
                    data = json.load(f)
                print("✓ Successfully loaded JSON data")
                
                if 'data' in data and isinstance(data['data'], list) and len(data['data']) > 0:
                    print(f"✓ Contains {len(data['data'])} data points")
                    first_point = data['data'][0]
                    if len(first_point) >= 6:
                        print(f"✓ Data points have correct format: {first_point[:6]}")
                    else:
                        print(f"✗ Data points have wrong format: {first_point}")
                else:
                    print("✗ No valid data array found")
                    
            except Exception as e:
                print(f"✗ Error loading file: {e}")
        else:
            print("✗ No trajectory files found")
    else:
        print("✗ Data directory does not exist")

if __name__ == "__main__":
    print("Testing Trajectory Visualization Logic")
    print("=" * 50)
    
    # Run all tests
    test_get_parcel_color_families()
    test_speed_to_color_intensity() 
    test_filename_parsing()
    test_trajectory_data_structure()
    verify_data_files()
    
    print("\n" + "=" * 50)
    print("Logic tests completed! The trajectory visualization code should work correctly.")
    print("Note: Actual visualization requires matplotlib and numpy to be properly installed.")