lap_position.py

import json
import os

import fastf1
import numpy as np
import utils
from fastf1.ergast import Ergast
import pandas as pd
fastf1.Cache.enable_cache("cache")
YEAR = 2020


events = [
    'Austrian Grand Prix', 
'Styrian Grand Prix', 
    'Hungarian Grand Prix',
    'British Grand Prix', 
    '70th Anniversary Grand Prix', 
    'Spanish Grand Prix', 'Belgian Grand Prix', 
    'Italian Grand Prix', 'Tuscan Grand Prix', 
    'Russian Grand Prix',
    'Eifel Grand Prix', 
    'Portuguese Grand Prix', 'Emilia Romagna Grand Prix', 
    'Turkish Grand Prix', 
    'Abu Dhabi Grand Prix',
    'Bahrain Grand Prix', 'Sakhir Grand Prix',   
]

def lap_chart(
    year: int,
    event: str | int,
    session: str,
) -> any:
    ergast = Ergast()

    race_names_df = ergast.get_race_schedule(season=year, result_type="pandas")
    event_number = race_names_df[race_names_df["raceName"] == event]["round"].values[0]
    drivers_df = ergast.get_driver_info(
        season=year, round=event_number, result_type="pandas"
    )
    laptimes_df = ergast.get_lap_times(
        season=year, round=event_number, result_type="pandas", limit=2000
    ).content[0]
    laptimes_df = pd.merge(laptimes_df, drivers_df, how="left", on="driverId")

    results_df = ergast.get_race_results(
        season=year, round=event_number, result_type="pandas"
    ).content[0]
    results_df = results_df[["driverCode", "constructorName"]]

    # merge results_df on laptime_df
    laptimes_df = pd.merge(laptimes_df, results_df, how="left", on="driverCode")

    team_colors = utils.team_colors(year)
    # add team_colors to laptimes_df
    laptimes_df["fill"] = laptimes_df["constructorName"].map(team_colors)

    #  rename number as x and position as y
    laptimes_df.rename(
        columns={"number": "x", "position": "y", "driverCode": "id"}, inplace=True
    )

    lap_chart_data = []

    for driver in laptimes_df["id"].unique():
        data = laptimes_df[laptimes_df["id"] == driver]
        fill = data["fill"].values[0]
        data = data[["x", "y"]]
        data_dict = data.to_dict(orient="records")
        driver_dict = {"id": driver, "fill": fill, "data": data_dict}
        # add this to all_data
        lap_chart_data.append(driver_dict)

    lap_chart_dict = {"lapChartData": lap_chart_data}

    return lap_chart_dict


# Your list of events
events_list = events

# Loop through each event
for event in events_list:
    session = "Race"

    lap_position_dict = lap_chart(YEAR, event, session)

    # Specify the file path where you want to save the JSON data
    file_path = f"{event}/{session}/lap_position.json"

    # Save the dictionary to a JSON file
    with open(file_path, "w") as json_file:
        json.dump(lap_position_dict, json_file)

    print(f"Dictionary saved to {file_path}")