F1Predictive_analysis.py

from pathlib import Path
import sys

import fastf1
from fastf1 import plotting
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures
from sklearn.metrics import r2_score
from sklearn.pipeline import make_pipeline
import warnings
warnings.filterwarnings("ignore")

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")

CACHE_DIR = Path(__file__).resolve().parent / "f1_cache"
CACHE_DIR.mkdir(exist_ok=True)
fastf1.Cache.enable_cache(str(CACHE_DIR))


# ─────────────────────────────────────────────
#  SESSION LOADING
# ─────────────────────────────────────────────

def session_data(year, gp):
    try:
        session = fastf1.get_session(year, gp, "R")
        session.load()
        return session
    except KeyboardInterrupt:
        print("\nSession download cancelled.")
        return None
    except Exception as e:
        print(f"\nError loading the {year} {gp} race: {e}")
        print("Check the season/race name and your internet connection, then try again.")
        return None


# ─────────────────────────────────────────────
#  DRIVER LISTING
# ─────────────────────────────────────────────

def racing_drivers(session):
    laps = session.laps
    race_drivers = laps["Driver"].unique()
    print("\nAvailable drivers in this race session:")
    for code in sorted(race_drivers):
        print(f"  - {code}")
    return race_drivers


# ─────────────────────────────────────────────
#  LAP TIME COMPARISON
# ─────────────────────────────────────────────

def compare_lap(driver1, driver2, session):
    laps = session.laps
    d1_laps = laps.pick_drivers(driver1).pick_quicklaps()
    d2_laps = laps.pick_drivers(driver2).pick_quicklaps()

    if d1_laps.empty or d2_laps.empty:
        missing = driver1 if d1_laps.empty else driver2
        raise ValueError(f"No valid quick laps were found for driver '{missing}'.")

    fastest1 = d1_laps.pick_fastest()
    fastest2 = d2_laps.pick_fastest()

    t1 = fastest1["LapTime"]
    t2 = fastest2["LapTime"]
    diff = (t1 - t2).total_seconds()

    avg1 = d1_laps["LapTime"].dt.total_seconds().mean()
    avg2 = d2_laps["LapTime"].dt.total_seconds().mean()

    print(f"\n{'─'*45}")
    print(f"  LAP TIME COMPARISON: {driver1} vs {driver2}")
    print(f"{'─'*45}")
    print(f"  {driver1} fastest : {t1}  |  avg: {avg1:.3f}s")
    print(f"  {driver2} fastest : {t2}  |  avg: {avg2:.3f}s")

    if diff < 0:
        print(f"\n  ✓ {driver1} was faster by {-diff:.3f} seconds")
    else:
        print(f"\n  ✓ {driver2} was faster by {diff:.3f} seconds")
    print(f"{'─'*45}\n")

    return t1, t2, diff


# ─────────────────────────────────────────────
#  PREDICTIVE MODEL
# ─────────────────────────────────────────────

def build_regression_model(lap_series, time_series, degree=2):
    """
    Fits a polynomial regression model to lap time data.
    Returns: model, r2_score, predicted array
    """
    X = lap_series.values.reshape(-1, 1)
    y = time_series.values

    model = make_pipeline(PolynomialFeatures(degree), LinearRegression())
    model.fit(X, y)
    y_pred = model.predict(X)
    r2 = r2_score(y, y_pred)

    return model, r2, y_pred


def predict_lap_time(model, lap_number):
    """Predicts lap time for a given lap number using the trained model."""
    X = np.array([[lap_number]])
    predicted = model.predict(X)[0]
    return predicted


def tire_degradation_rate(lap_numbers, lap_times):
    """
    Estimates tire degradation as the slope of a simple linear regression
    on the lap time series. A positive slope = degradation.
    """
    X = np.array(lap_numbers).reshape(-1, 1)
    y = np.array(lap_times)
    model = LinearRegression().fit(X, y)
    return model.coef_[0]  # seconds per lap


def print_predictive_summary(driver1, driver2, model1, model2, r2_1, r2_2,
                              d1_laps, d2_laps):
    x1 = d1_laps["LapNumber"].values
    y1 = d1_laps["LapTime"].dt.total_seconds().values
    x2 = d2_laps["LapNumber"].values
    y2 = d2_laps["LapTime"].dt.total_seconds().values

    deg1 = tire_degradation_rate(x1, y1)
    deg2 = tire_degradation_rate(x2, y2)

    print(f"\n{'─'*45}")
    print("  PREDICTIVE ANALYSIS SUMMARY")
    print(f"{'─'*45}")
    print(f"  {driver1} regression R²  : {r2_1:.4f}")
    print(f"  {driver2} regression R²  : {r2_2:.4f}")
    print(f"  {driver1} tire degradation : {deg1:+.4f} s/lap")
    print(f"  {driver2} tire degradation : {deg2:+.4f} s/lap")

    better = driver1 if deg1 < deg2 else driver2
    print(f"\n  → {better} shows better tire management (lower degradation)")
    print(f"{'─'*45}\n")

    return deg1, deg2


def interactive_prediction(driver1, driver2, model1, model2):
    """Lets the user predict lap times at any lap number."""
    print("\n  PREDICTION MODE")
    print("  Enter a lap number to predict times (or 0 to exit):\n")

    while True:
        try:
            lap = int(input("  Lap number: "))
            if lap == 0:
                break
            p1 = predict_lap_time(model1, lap)
            p2 = predict_lap_time(model2, lap)
            diff = abs(p1 - p2)
            faster = driver1 if p1 < p2 else driver2
            print(f"    {driver1} predicted: {p1:.3f}s")
            print(f"    {driver2} predicted: {p2:.3f}s")
            print(f"    → {faster} faster by {diff:.3f}s at lap {lap}\n")
        except ValueError:
            print("  Please enter a valid integer.")


# ─────────────────────────────────────────────
#  VISUALIZATIONS
# ─────────────────────────────────────────────

def plot_full_analysis(driver1, driver2, session,
                       model1, model2, r2_1, r2_2,
                       deg1, deg2):
    plotting.setup_mpl(misc_mpl_mods=False)
    laps = session.laps

    d1 = laps.pick_drivers(driver1).pick_quicklaps()
    d2 = laps.pick_drivers(driver2).pick_quicklaps()

    x1 = d1["LapNumber"].values
    y1 = d1["LapTime"].dt.total_seconds().values
    x2 = d2["LapNumber"].values
    y2 = d2["LapTime"].dt.total_seconds().values

    lap_range = np.linspace(
        min(x1.min(), x2.min()),
        max(x1.max(), x2.max()),
        200
    ).reshape(-1, 1)

    pred1_range = model1.predict(lap_range)
    pred2_range = model2.predict(lap_range)

    # ── Colors ──────────────────────────────
    try:
        c1 = plotting.get_driver_color(driver1, session)
        c2 = plotting.get_driver_color(driver2, session)
    except Exception:
        c1, c2 = "#378add", "#e24b4a"

    fig = plt.figure(figsize=(16, 12), facecolor="#0d0d0d")
    fig.suptitle(
        f"F1 Predictive Analysis — {driver1} vs {driver2}",
        fontsize=16, fontweight="bold", color="white", y=0.97
    )
    gs = gridspec.GridSpec(2, 2, figure=fig, hspace=0.42, wspace=0.32)

    # ── Plot 1: Lap times + regression ──────
    ax1 = fig.add_subplot(gs[0, :])
    ax1.set_facecolor("#1a1a1a")
    ax1.scatter(x1, y1, color=c1, alpha=0.55, s=18, zorder=3, label=f"{driver1} laps")
    ax1.scatter(x2, y2, color=c2, alpha=0.55, s=18, zorder=3, label=f"{driver2} laps")
    ax1.plot(lap_range, pred1_range, color=c1, lw=2, linestyle="--",
             label=f"{driver1} regression (R²={r2_1:.3f})", zorder=4)
    ax1.plot(lap_range, pred2_range, color=c2, lw=2, linestyle="--",
             label=f"{driver2} regression (R²={r2_2:.3f})", zorder=4)
    ax1.set_xlabel("Lap Number", color="gray", fontsize=11)
    ax1.set_ylabel("Lap Time (seconds)", color="gray", fontsize=11)
    ax1.set_title("Lap Time Trend + Polynomial Regression Model", color="white", fontsize=12)
    ax1.tick_params(colors="gray")
    ax1.spines[["top", "right"]].set_visible(False)
    ax1.spines[["bottom", "left"]].set_color("#333")
    ax1.grid(alpha=0.12, color="white")
    ax1.legend(fontsize=10, facecolor="#111", labelcolor="white", framealpha=0.8)

    # ── Plot 2: Residuals ────────────────────
    ax2 = fig.add_subplot(gs[1, 0])
    ax2.set_facecolor("#1a1a1a")
    res1 = y1 - model1.predict(x1.reshape(-1, 1))
    res2 = y2 - model2.predict(x2.reshape(-1, 1))
    ax2.scatter(x1, res1, color=c1, alpha=0.6, s=16, label=driver1)
    ax2.scatter(x2, res2, color=c2, alpha=0.6, s=16, label=driver2)
    ax2.axhline(0, color="white", lw=0.8, linestyle="-", alpha=0.3)
    ax2.set_xlabel("Lap Number", color="gray", fontsize=10)
    ax2.set_ylabel("Residual (s)", color="gray", fontsize=10)
    ax2.set_title("Model Residuals", color="white", fontsize=11)
    ax2.tick_params(colors="gray")
    ax2.spines[["top", "right"]].set_visible(False)
    ax2.spines[["bottom", "left"]].set_color("#333")
    ax2.grid(alpha=0.12, color="white")
    ax2.legend(fontsize=9, facecolor="#111", labelcolor="white", framealpha=0.8)

    # ── Plot 3: Lap time distribution ───────
    ax3 = fig.add_subplot(gs[1, 1])
    ax3.set_facecolor("#1a1a1a")
    bins = np.linspace(min(y1.min(), y2.min()) - 0.3,
                       max(y1.max(), y2.max()) + 0.3, 20)
    ax3.hist(y1, bins=bins, color=c1, alpha=0.65, label=driver1, edgecolor="none")
    ax3.hist(y2, bins=bins, color=c2, alpha=0.65, label=driver2, edgecolor="none")
    ax3.axvline(y1.mean(), color=c1, lw=1.5, linestyle=":", alpha=0.9)
    ax3.axvline(y2.mean(), color=c2, lw=1.5, linestyle=":", alpha=0.9)
    ax3.set_xlabel("Lap Time (s)", color="gray", fontsize=10)
    ax3.set_ylabel("Count", color="gray", fontsize=10)
    ax3.set_title("Lap Time Distribution", color="white", fontsize=11)
    ax3.tick_params(colors="gray")
    ax3.spines[["top", "right"]].set_visible(False)
    ax3.spines[["bottom", "left"]].set_color("#333")
    ax3.grid(alpha=0.12, color="white")
    ax3.legend(fontsize=9, facecolor="#111", labelcolor="white", framealpha=0.8)

    # ── Degradation annotation ───────────────
    fig.text(0.5, 0.005,
             f"Tire degradation — {driver1}: {deg1:+.4f}s/lap  |  {driver2}: {deg2:+.4f}s/lap",
             ha="center", fontsize=10, color="gray")

    output_path = Path(__file__).resolve().parent / "f1_predictive_analysis.png"
    plt.savefig(output_path, dpi=150, bbox_inches="tight",
                facecolor="#0d0d0d")
    print(f"\n  Chart saved successfully:\n  {output_path}")
    plt.show()
    plt.close(fig)


# ─────────────────────────────────────────────
#  MAIN
# ─────────────────────────────────────────────

def main():
    print("\n" + "═" * 47)
    print("  🏎️  F1 PREDICTIVE ANALYSIS SYSTEM")
    print("═" * 47)

    try:
        year = int(input("\n  Season year (e.g. 2023): "))
    except ValueError:
        print("  Season year must be a number, for example 2023.")
        return

    gp = input("  Grand Prix (e.g. Monaco): ").strip()
    if not gp:
        print("  Grand Prix cannot be empty.")
        return

    session = session_data(year, gp)
    if session is None:
        return

    available_drivers = set(racing_drivers(session))

    driver1 = input("\n  First driver code  (e.g. VER): ").strip().upper()
    driver2 = input("  Second driver code (e.g. LEC): ").strip().upper()
    invalid = [
        driver for driver in (driver1, driver2)
        if driver not in available_drivers
    ]
    if invalid:
        print(f"  Invalid driver code(s): {', '.join(invalid)}")
        print("  Choose codes from the available-driver list above.")
        return
    if driver1 == driver2:
        print("  Please choose two different drivers.")
        return

    try:
        # Basic comparison
        compare_lap(driver1, driver2, session)

        # Prepare data
        laps = session.laps
        d1_laps = laps.pick_drivers(driver1).pick_quicklaps()
        d2_laps = laps.pick_drivers(driver2).pick_quicklaps()

        x1 = d1_laps["LapNumber"]
        y1 = d1_laps["LapTime"].dt.total_seconds()
        x2 = d2_laps["LapNumber"]
        y2 = d2_laps["LapTime"].dt.total_seconds()

        # Build regression models (degree=2 polynomial)
        model1, r2_1, _ = build_regression_model(x1, y1, degree=2)
        model2, r2_2, _ = build_regression_model(x2, y2, degree=2)

        # Predictive summary
        deg1, deg2 = print_predictive_summary(
            driver1, driver2, model1, model2, r2_1, r2_2, d1_laps, d2_laps
        )

        # Generate and save the chart before entering interactive prediction mode.
        plot_full_analysis(
            driver1, driver2, session,
            model1, model2, r2_1, r2_2,
            deg1, deg2
        )

        # Interactive prediction
        interactive_prediction(driver1, driver2, model1, model2)
    except (ValueError, KeyError) as e:
        print(f"\nUnable to analyze the selected drivers: {e}")


if __name__ == "__main__":
    main()