f1-circuits/coordinate_generator.py

import fastf1
import numpy as np
import matplotlib.pyplot as plt
import os
import json
import re
from typing import List, Dict, Any, Optional, Tuple

# Import for type hints
import fastf1.core

def rotate(xy, *, angle):
    """Rotate coordinates by the given angle."""
    rot_mat = np.array([[np.cos(angle), np.sin(angle)],
                        [-np.sin(angle), np.cos(angle)]])
    return np.matmul(xy, rot_mat)

def get_track_coordinates(session: fastf1.core.Session):
    """Extract track coordinates from session data."""
    try:
        # Get the fastest lap and position data
        lap = session.laps.pick_fastest()
        pos = lap.get_pos_data()
        
        # Get circuit information
        circuit_info = session.get_circuit_info()
        
        # Get track coordinates
        track = pos.loc[:, ('X', 'Y')].to_numpy()
        
        # Convert the rotation angle from degrees to radian
        track_angle = getattr(circuit_info, 'rotation', 0) / 180 * np.pi
        
        # Rotate the track coordinates
        rotated_track = rotate(track, angle=track_angle)
        
        # Create circuit_info dictionary with safe attribute access
        circuit_info_dict = {
            'name': session.event['EventName'],
            'country': session.event['Country'],
            'city': session.event['Location'],
            'length': getattr(circuit_info, 'length', None),
            'rotation': getattr(circuit_info, 'rotation', None),
        }
        
        # Create event_info dictionary
        event_info_dict = {
            'name': session.event['EventName'],
            'year': session.event['EventDate'].year,
            'official_name': session.event.get('OfficialEventName', session.event['EventName']),
            'location': session.event['Location']
        }
        
        return {
            'coordinates': rotated_track.tolist(),
            'circuit_info': circuit_info_dict,
            'event_info': event_info_dict
        }
    except Exception as e:
        print(f"Error extracting track coordinates: {str(e)}")
        import traceback
        traceback.print_exc()
        return None

def create_geojson(track_data):
    """Create a GeoJSON representation of the track."""
    # Extract data from track_data
    coordinates = track_data['coordinates']
    circuit_info = track_data['circuit_info']
    event_info = track_data['event_info']
    
    # Create a simplified ID based on country and year
    country_code = get_country_code(circuit_info['country'])
    year = event_info['year']
    circuit_id = f"{country_code.lower()}-{year}"
    
    # Ensure length is an integer if available, or None if not
    circuit_length = None
    if circuit_info.get('length'):
        try:
            circuit_length = int(circuit_info['length'])
        except (ValueError, TypeError):
            pass
    
    # Create GeoJSON structure
    geojson = {
        "type": "FeatureCollection",
        "name": circuit_id,
        "features": [
            {
                "type": "Feature",
                "properties": {
                    "id": circuit_id,
                    "Location": event_info['location'],
                    "Name": circuit_info['name'],
                    "opened": None,  # Not available from API
                    "firstgp": None, # Not available from API
                    "length": circuit_length,
                    "altitude": None  # Not available from API
                },
                "geometry": {
                    "type": "LineString",
                    "coordinates": coordinates
                }
            }
        ]
    }
    
    # Calculate bounding box if coordinates are available
    if coordinates and len(coordinates) > 0:
        x_coords = [coord[0] for coord in coordinates]
        y_coords = [coord[1] for coord in coordinates]
        bbox = [min(x_coords), min(y_coords), max(x_coords), max(y_coords)]
        
        # Add bounding box to GeoJSON
        geojson["bbox"] = bbox
        geojson["features"][0]["bbox"] = bbox
    
    return geojson

def get_country_code(country_name):
    """Get a two-letter country code from a country name."""
    country_codes = {
        "Australia": "AU",
        "Austria": "AT",
        "Azerbaijan": "AZ",
        "Bahrain": "BH",
        "Belgium": "BE",
        "Brazil": "BR",
        "Canada": "CA",
        "China": "CN",
        "Denmark": "DK",
        "France": "FR",
        "Germany": "DE",
        "Hungary": "HU",
        "Italy": "IT",
        "Japan": "JP",
        "Mexico": "MX",
        "Monaco": "MC",
        "Netherlands": "NL",
        "Portugal": "PT",
        "Qatar": "QA",
        "Russia": "RU",
        "Saudi Arabia": "SA",
        "Singapore": "SG",
        "Spain": "ES",
        "Sweden": "SE",
        "Switzerland": "CH",
        "United Arab Emirates": "AE",
        "United Kingdom": "GB",
        "United States": "US",
        "USA": "US",
        "Abu Dhabi": "AE",  # Special case for events
        "Great Britain": "GB",
        "UK": "GB",
        "Emilia Romagna": "IT",  # Special case for events
        "Styria": "AT",  # Special case for events
        "Miami": "US",
        "Las Vegas": "US"
    }
    
    return country_codes.get(country_name, "XX")

def sanitize_filename(name):
    """Convert a string to a valid filename."""
    # Replace special characters with spaces
    name = re.sub(r'[\\/*?:"<>|]', ' ', name)
    # Replace multiple spaces with single space
    name = re.sub(r'\s+', ' ', name).strip()
    return name

def save_svg(coordinates, output_path):
    """Generate and save an SVG representation of the circuit."""
    if not coordinates:
        print("No coordinates available to create SVG")
        return False
    
    # Extract x and y coordinates
    x_coords = [coord[0] for coord in coordinates]
    y_coords = [coord[1] for coord in coordinates]
    
    # Normalize coordinates to fit in SVG
    min_x, max_x = min(x_coords), max(x_coords)
    min_y, max_y = min(y_coords), max(y_coords)
    
    # Add some padding
    x_padding = (max_x - min_x) * 0.05
    y_padding = (max_y - min_y) * 0.05
    
    min_x -= x_padding
    max_x += x_padding
    min_y -= y_padding
    max_y += y_padding
    
    # SVG dimensions
    svg_width = 800
    svg_height = int(svg_width * (max_y - min_y) / (max_x - min_x))
    
    # Create SVG header
    svg = [
        f'<?xml version="1.0" encoding="UTF-8" standalone="no"?>',
        f'<svg width="{svg_width}" height="{svg_height}" viewBox="0 0 {svg_width} {svg_height}" xmlns="http://www.w3.org/2000/svg">',
    ]
    
    # Create path data
    path_data = []
    for i, (x, y) in enumerate(coordinates):
        # Convert coordinates to SVG coordinates
        svg_x = svg_width * (x - min_x) / (max_x - min_x)
        # Flip y-axis because SVG has origin at top-left
        svg_y = svg_height * (1 - (y - min_y) / (max_y - min_y))
        
        if i == 0:
            path_data.append(f"M{svg_x:.2f},{svg_y:.2f}")
        else:
            path_data.append(f"L{svg_x:.2f},{svg_y:.2f}")
    
    # Add path to SVG
    svg.append(f'  <path d="{" ".join(path_data)}" fill="none" stroke="black" stroke-width="2"/>')
    svg.append('</svg>')
    
    # Write SVG to file
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write('\n'.join(svg))
    
    return True

def save_png(coordinates, output_path):
    """Generate and save a PNG representation of the circuit."""
    if not coordinates:
        print("No coordinates available to create PNG")
        return False
    
    # Extract x and y coordinates
    x_coords = [coord[0] for coord in coordinates]
    y_coords = [coord[1] for coord in coordinates]
    
    # Create figure with transparent background
    fig = plt.figure(figsize=(10, 8), facecolor="none")
    ax = fig.add_subplot(111)
    
    # Plot the circuit with black line
    ax.plot(x_coords, y_coords, 'k-', linewidth=2)
    
    # Remove grid, axes, and border
    ax.grid(False)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['bottom'].set_visible(False)
    ax.spines['left'].set_visible(False)
    
    # Keep aspect ratio equal
    ax.set_aspect('equal')
    
    # Add padding around the circuit
    ax.margins(0.1)
    
    # Save with transparent background
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    fig.savefig(output_path, transparent=True, bbox_inches='tight', pad_inches=0.1)
    plt.close(fig)
    
    return True

# Additional utility functions for interactive mode
def get_available_seasons():
    """Get a list of available F1 seasons."""
    from datetime import datetime
    current_year = datetime.now().year
    # F1 API has data from 1950 to current year
    return list(range(1950, current_year + 1))

def get_events_for_season(year):
    """Get all events for a specific season."""
    try:
        schedule = fastf1.get_event_schedule(year)
        return schedule
    except Exception as e:
        print(f"Error fetching events for {year}: {str(e)}")
        return None