import os
from shiny import App, ui, render, reactive
import fastf1 as ff1
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib import cm
import numpy as np
import shiny.experimental as x
import shinyswatch

# Define cache folder path
cache_path = os.getcwd() + "/cache"
print(f"Cache path: {cache_path}")

ff1.Cache.enable_cache(cache_path)
#ff1.Cache.offline_mode(enabled=True)

app_ui = ui.page_fluid(
    shinyswatch.theme.minty(),
    ui.layout_sidebar(
        ui.panel_sidebar(
            ui.input_select(
            "track", label="Select track:",
            choices=["Austria", "Hungary", "Spain", "Bahrain"],
            selected = "Austria"),
            ui.input_radio_buttons(
                    "session_type", "Session type:", 
                    choices = {"R": "Race", "Q": "Qualification"}, 
                    selected = "R"),
            ui.input_radio_buttons(
                    "year", "Year:", 
                    choices = [2023, 2022], 
                    selected = 2023),
            ),

        ui.panel_main(
            ui.h3("Gear usage in fastest lap"),
            x.ui.output_plot("gear", fill = True),
            ui.output_text("fastest_driver")
        ),
    ),
)   

def server(input, output, session):
    @reactive.Calc
    def get_data():
        f1_session = ff1.get_session(int(input.year()), input.track(), input.session_type())
        f1_session.load()

        lap = f1_session.laps.pick_fastest()
        tel = lap.get_telemetry()
        driver = lap['Driver']

        #converting data to numpy data tables
        x = np.array(tel['X'].values)
        y = np.array(tel['Y'].values)

        points = np.array([x, y]).T.reshape(-1, 1, 2)
        segments = np.concatenate([points[:-1], points[1:]], axis=1)
        gear = tel['nGear'].to_numpy().astype(float)
        return segments, gear, driver

    @output
    @render.plot
    def gear():
        segments, gear, driver = get_data()
        
        cmap = cm.get_cmap('Paired')
        lc_comp = LineCollection(segments, norm=plt.Normalize(1, cmap.N+1), cmap=cmap)
        lc_comp.set_array(gear)
        lc_comp.set_linewidth(4)

        plt.gca().add_collection(lc_comp)
        plt.axis('equal')
        plt.tick_params(labelleft=False, left=False, labelbottom=False, bottom=False)

        cbar = plt.colorbar(mappable=lc_comp, label="Gear", boundaries=np.arange(1, 10))
        cbar.set_ticks(np.arange(1.5, 9.5))
        cbar.set_ticklabels(np.arange(1, 9))
            
        plt

    @output
    @render.text
    def fastest_driver():
        segments, gear, driver = get_data()
        return f"The driver of the fastest lap is: {driver}"


app = App(app_ui, server)