tools.py

"""
Here we implement the MCP tools

Tools to implement
- session result
- driver info
- championship standings
- driver standings
- constructor standings
- compare drivers

"""

import json
import fastf1
from PIL import Image
from typing import Union
from fastf1.core import Session

# Local modules
from utils import parser_utils, track_utils
from utils.constants import AVAILABLE_SESSION_TYPES

# Custom types
gp = Union[str, int]
session_type = Union[str, int, None]

def get_session(year: int, round: gp, session_type: session_type) -> Session:
    """ Returns a session specified by year, round and session type """

    # Check if session type is valid
    if isinstance(session_type, str):
        if session_type.lower() not in AVAILABLE_SESSION_TYPES:
            return f"Session type {session_type} is not available. Supported session types: {AVAILABLE_SESSION_TYPES}"

    return fastf1.get_session(year, round, session_type)

def get_season_calendar(year: int) -> str:
    """ Returns the season calendar for the given year """
    season_calendar = fastf1.get_event_schedule(year)
    return parser_utils.parse_season_calendar(season_calendar)

def get_event_info(year: int, round: gp, format: str) -> str:
    """ Returns the event info for the specified session """

    event = fastf1.get_session(year, round, "race").event # Event object is the same for all sessions, so hardcode "race"
    if format == "human":
        data_interval = f"{event['Session1DateUtc'].date()} - {event['Session5DateUtc'].date()}"
        event_string = f"Round {event['RoundNumber']} : {event['EventName']} - {event['Location']}, {event['Country']} ({data_interval})"
        return event_string
    elif format == "LLM":
        return parser_utils.parse_event_info(event)


def get_constructor_standings(year: int) -> str:
    pass

def get_driver_standings(year: int) -> str: 
    pass

def driver_championship_standings(year: int, driver_name: str) -> str:

    with open("assets/driver_abbreviations.json") as f:
        driver_abbreviations = json.load(f)
        driver_abbreviation = driver_abbreviations[driver_name]
    ergast = fastf1.ergast.Ergast()
    driver_standings = ergast.get_driver_standings(year).content[0]
    driver_standing = driver_standings[["position", "points", "wins", "driverCode"]].reset_index(drop=True)
    driver_standing = driver_standing[driver_standing["driverCode"] == driver_abbreviation]
    suffix = "st" if driver_standing['position'].iloc[0] == 1 else "nd" if driver_standing['position'].iloc[0] == 2 else "rd" if driver_standing['position'].iloc[0] == 3 else "th"
    standings_string = f"{driver_name} is {driver_standing['position'].iloc[0]}{suffix} with {driver_standing['points'].iloc[0]} points and {driver_standing['wins'].iloc[0]} wins"
    return standings_string
    

def constructor_championship_standings(year: int, constructor_name: str) -> str:

    team_mapping = {
    "McLaren": "McLaren",
    "Ferrari": "Ferrari",
    "Red Bull Racing": "Red Bull",
    "Mercedes": "Mercedes",
    "Aston Martin": "Aston Martin",
    "Alpine": "Alpine F1 Team",
    "Haas": "Haas F1 Team",
    "Racing Bulls": "RB F1 Team",
    "Williams": "Williams",
    "Kick Sauber": "Sauber"
    }

    ergast = fastf1.ergast.Ergast()
    constructor_standings = ergast.get_constructor_standings(year).content[0]
    constructor_standing = constructor_standings[["position", "points", "wins", "constructorName"]].reset_index(drop=True)
    mapped_name = team_mapping[constructor_name]
    constructor_standing = constructor_standing[constructor_standing["constructorName"] == mapped_name]
    suffix = "st" if constructor_standing['position'].iloc[0] == 1 else "nd" if constructor_standing['position'].iloc[0] == 2 else "rd" if constructor_standing['position'].iloc[0] == 3 else "th"
    standings_string = f"{constructor_name} are {constructor_standing['position'].iloc[0]}{suffix} with {constructor_standing['points'].iloc[0]} points and {constructor_standing['wins'].iloc[0]} wins"
    return standings_string



def track_visualization(year: int, round: gp, visualization_type: str, driver_name: str) -> Image:

    session = get_session(year, round, "race")
    session.load()

    if visualization_type == "speed":
        return track_utils.create_track_speed_visualization(session, driver_name)
    elif visualization_type == "corners":
        return track_utils.create_track_corners_visualization(session)
    elif visualization_type == "gear":
        return track_utils.create_track_gear_visualization(session)


if __name__ == "__main__":
    session = get_session(2024, 1, "fp1")
    session.load()