import fastf1
import numpy as np
import matplotlib.pyplot as plt
import os
import json
from datetime import datetime
import re
import fastf1.core as f1

# Enable cache to speed up subsequent data loading
fastf1.Cache.enable_cache('cache')

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: f1.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()
        
        # Print available attributes for debugging
        print("\nAvailable circuit_info attributes:")
        for attr in dir(circuit_info):
            if not attr.startswith('_'):  # Skip private attributes
                try:
                    value = getattr(circuit_info, attr)
                    print(f"  {attr}: {value}")
                except Exception as e:
                    print(f"  {attr}: Error accessing ({str(e)})")
        
        # Get track coordinates
        track = pos.loc[:, ('X', 'Y')].to_numpy()
        
        # Convert the rotation angle from degrees to radian
        track_angle = circuit_info.rotation / 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

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

def prompt_season_selection():
    """Prompt the user to select a season."""
    seasons = get_available_seasons()
    
    print("\nAvailable F1 Seasons:")
    
    # Display seasons in rows of 10
    for i in range(0, len(seasons), 10):
        row = seasons[i:i+10]
        print(" ".join(f"{year:<6}" for year in row))
    
    while True:
        try:
            selection = input("\nEnter season year (e.g., 2023): ")
            year = int(selection)
            if year in seasons:
                return year
            else:
                print(f"Invalid season: {year}. Please choose a year between 1950 and {seasons[-1]}.")
        except ValueError:
            print("Please enter a valid year (e.g., 2023).")

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

def prompt_event_selection(events):
    """Prompt the user to select an event from the schedule."""
    if events is None or len(events) == 0:
        print("No events available for the selected season.")
        return None
    
    print("\nAvailable Grands Prix:")
    for i, (_, event) in enumerate(events.iterrows(), 1):
        event_name = event['EventName']
        location = event['Location']
        date = event['EventDate'].strftime('%Y-%m-%d')
        print(f"{i:2}. {event_name:<35} {location:<20} {date}")
    
    while True:
        try:
            selection = input("\nEnter the number of the Grand Prix: ")
            idx = int(selection) - 1
            if 0 <= idx < len(events):
                return events.iloc[idx]
            else:
                print(f"Invalid selection. Please choose a number between 1 and {len(events)}.")
        except ValueError:
            print("Please enter a valid number.")

def create_output_directory(country, location, circuit_name, year):
    """Create the appropriate directory structure for saving files."""
    # Sanitize names for directory structure
    country = sanitize_filename(country)
    location = sanitize_filename(location)
    circuit_name = sanitize_filename(circuit_name)
    
    # Create directory path
    base_dir = os.path.join("circuits", country, location)
    os.makedirs(base_dir, exist_ok=True)
    
    # Create filename base (without extension)
    filename_base = f"{circuit_name} - {year}-"
    
    return base_dir, filename_base

def main():
    try:
        print("F1 Circuit Coordinate Generator")
        print("===============================")
        print("This tool allows you to download coordinates for F1 circuits and save them as GeoJSON, SVG, and PNG files.")
        
        # Step 1: Select a season
        year = prompt_season_selection()
        print(f"\nSelected season: {year}")
        
        # Step 2: Get events for the selected season
        print(f"\nFetching events for {year}...")
        events = get_events_for_season(year)
        
        if events is None or len(events) == 0:
            print(f"No events found for the {year} season. Exiting.")
            return
        
        # Step 3: Select an event
        event = prompt_event_selection(events)
        if event is None:
            print("No event selected. Exiting.")
            return
        
        print(f"\nSelected event: {event['EventName']} at {event['Location']}")
        
        # Step 4: Download the session data (try qualifying first, then other sessions)
        print("\nDownloading session data (this may take a moment)...")
        session_types = ['Q', 'FP1', 'FP2', 'FP3', 'R']
        
        session = None
        for session_type in session_types:
            try:
                print(f"Trying to load {session_type} session...")
                session = fastf1.get_session(year, event['EventName'], session_type)
                session.load()
                if session and session.total_laps is not None and session.total_laps > 0:
                    print(f"Successfully loaded {session_type} session")
                    break
            except Exception as e:
                print(f"Could not load {session_type} session: {str(e)}")
        
        if session is None:
            print("Could not load any session data. Exiting.")
            return
        
        # Step 5: Extract track coordinates
        print("\nExtracting track coordinates...")
        track_data = get_track_coordinates(session)
        
        if not track_data:
            print("Failed to extract track coordinates. Exiting.")
            return
        
        # Step 6: Create GeoJSON data
        print("\nCreating GeoJSON data...")
        geojson_data = create_geojson(track_data)
        
        # Step 7: Determine output directory and filenames
        circuit_info = track_data['circuit_info']
        event_info = track_data['event_info']
        
        country = circuit_info['country']
        location = event_info['location']
        circuit_name = circuit_info['name']
        year = event_info['year']
        
        print(f"\nPreparing to save circuit: {circuit_name} ({location}, {country}) from {year}")
        
        base_dir, filename_base = create_output_directory(country, location, circuit_name, year)
        
        geojson_path = os.path.join(base_dir, f"{filename_base}.geo.json")
        svg_path = os.path.join(base_dir, f"{filename_base}.svg")
        png_path = os.path.join(base_dir, f"{filename_base}.png")
        
        # Step 8: Save files
        print("\nSaving files...")
        
        # Save GeoJSON
        with open(geojson_path, 'w', encoding='utf-8') as f:
            json.dump(geojson_data, f, indent=2)
        print(f"Saved GeoJSON to: {geojson_path}")
        
        # Save SVG
        if save_svg(track_data['coordinates'], svg_path):
            print(f"Saved SVG to: {svg_path}")
        
        # Save PNG
        if save_png(track_data['coordinates'], png_path):
            print(f"Saved PNG to: {png_path}")
        
        print("\nDone! Files have been saved successfully.")
        
    except Exception as e:
        print(f"\nAn error occurred: {str(e)}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()