from shiny import App, ui, render, reactive
from shinywidgets import output_widget, render_widget
import fastf1 as ff1
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib import cm
import numpy as np

app_ui = ui.page_fluid(
    ui.div(
        ui.input_select(
            "track", label="Track",
            choices=["Austria", "Hungary"]
        ),
        class_="d-flex gap-3"
    ),
    output_widget("my_widget")
)

def server(input, output, session):
    @reactive.Calc
    def get_data():
        ff1.Cache.enable_cache('.\cache')

        session_f1 = ff1.get_session(2023, input.track, 'R')
        session_f1.load()

        lap = session_f1.laps.pick_fastest()
        tel = lap.get_telemetry()

        #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 lap, tel, x,y,points,segments,gear

    @output
    @render_widget
    def my_widget(lap, tel, x,y,points,segments,gear):
        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)

        title = plt.suptitle(
                f"Fastest Lap Gear Shift Visualization\n"
                f"{lap['Driver']} - {session.event['EventName']} {session.event.year}"
                )

        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))
            
    return plt


app = App(app_ui, server)