# For reading PDF pages
import fitz

# For plotting images
import matplotlib.pyplot as plt

# To convert byte data to image objects
from PIL import Image

# For in-memory byte handling
from io import BytesIO

# System utilities for path management
import sys
import os
sys.path.append(os.path.abspath(".."))

# Custom utility to get driver names by number
from fixedData.drivers import getDriver

__all__ = ["showRaceStrategy"]

def showRaceStrategy(race, driverNumber, firstLap, lastLap):
    """
    Displays a grid of race strategy images (e.g., pit stop charts) for a given Formula 1 driver.

    Parameters:
        race (str): Name of the race (used as folder name where lap PDFs are stored).
        driverNumber (int): The car number of the F1 driver.
        firstLap (int): The first lap to include.
        lastLap (int): The last lap to include.
    """

    # Load images from PDF files
    images = []
    for i in range(firstLap, lastLap + 1):
        file = os.path.join(race, f"lap{i}.pdf")
        if os.path.exists(file):
            doc = fitz.open(file)
            page = doc.load_page(0)
            pix = page.get_pixmap()
            img = Image.open(BytesIO(pix.tobytes("png")))
            images.append((i, img))
            doc.close()

    # Plot images using subplots
    cols = 5
    rows = (len(images) + cols - 1) // cols

    fig, axes = plt.subplots(rows, cols, figsize=(15, rows * 2.5))
    axes = axes.flatten()

    for ax, (i, img) in zip(axes, images):
        ax.imshow(img)
        ax.set_title(f"Lap {i}")
        ax.axis("off")

    # Hide unused subplots
    for ax in axes[len(images):]:
        ax.axis("off")

    # Add figure title with driver and race info
    fig.suptitle(f"Pit stop strategy for {getDriver(driverNumber)} in {race}", fontsize=16, fontweight="bold")
    plt.tight_layout()

    # Save the final figure
    outputPath = os.path.join(race, f"fullStrategy.pdf")
    fig.savefig(outputPath, format="pdf")

    plt.show()