import os
from shiny import module, 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

# UI
@module.ui
def plot1_ui(label: str = 'plot1'):
    return ui.div(
        ui.output_plot('plot1')
    )

# Server
@module.server
def plot1_server(input, output, session):
    @reactive.Calc
    # Get required data for driver 1 based on selection
    def get_data_1():
        try:
            ui.notification_show("Data takes a couple seconds to load.", duration=3, type = 'default')
            
            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.driver1_select() == "Fastest driver":
                lap = f1_session.laps.pick_fastest()
            else:
                laps_driver = f1_session.laps.pick_driver(input.driver1_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)
            lap_time = lap['LapTime']
            return segments, gear, driver, lap_time

        except Exception:
            ui.notification_show("Data not available. Select another track or driver.", duration=10, type = 'error')

        @output
        @render.plot
        def gear_1():
            try:
                segments, gear, driver, lap_time = get_data_1()
                
                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

            except Exception:
                pass