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 shinyswatch

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

ff1.Cache.enable_cache(cache_path)

# Offline mode to prevent F1 API crashes on Hugging Face
ff1.Cache.offline_mode(enabled=True)

app_ui = ui.page_fluid(
    shinyswatch.theme.minty(),
    ui.panel_title("Gear usage in fastest lap"),
    ui.layout_sidebar(
        ui.panel_sidebar(
            ui.input_select(
                "track_select", "Select track:",
                choices = ["Austria", "Hungary", "Spain", "Bahrain"],
                selected = "Austria"
            ),
            ui.input_select(
                "driver_select", label="Select driver:",
                choices = ["Fastest driver"],
                selected = "Fastest driver"
            ),
            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"
            ),
            width=2
        ),
        
        ui.panel_main(
            ui.output_plot("gear"),
            ui.output_text("fastest_driver")
        ),
    ),
)   

def server(input, output, session):
    @reactive.Effect()
    def _():
        if input.year() == "2023":
            driver_options = ['Fastest driver', 'RUS', 'VER', 'HAM']
        elif input.year() == "2022":
            driver_options = ['Fastest driver', 'LEC', 'NOR', 'GAS', 'VER']
            
        ui.update_select("driver_select",
            label="Select driver:",
            choices=driver_options,
            selected=input.driver_select()
        )
    
    @reactive.Calc
    def get_data():
        f1_session = ff1.get_session(int(input.year()), input.track_select(), input.session_type())
        f1_session.load()

        # Check if user input == fastest driver
        if input.driver_select() == "Fastest driver":
            lap = f1_session.laps.pick_fastest()
        else:
            laps_driver = f1_session.laps.pick_driver(input.driver_select())
            lap = laps_driver.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.text
    def fastest_driver():
        segments, gear, driver = get_data()
        #print(f"The driver of the fastest lap this session is: {driver}")
        return f"Fastest lap this session of: {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


app = App(app_ui, server)