import fastf1 as ff1
from fastf1 import utils
from fastf1 import plotting
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import rcParams
rcParams['font.family']='serif'

plotting.setup_mpl(color_scheme='None')
ff1.Cache.enable_cache('fastf1cache')

year = 2021
GP = 'Saudi Arabia Grand Prix'
event = 'R'
lap = '37'

race = ff1.get_session(year, GP, event)
laps = race.load_laps(with_telemetry=True)
lap  = laps.loc[laps['LapNumber']==37]
ver  = lap.pick_driver('VER').get_telemetry()

verd = np.asarray(ver['Distance'])
vers = np.asarray(ver['Speed'])
vert = np.asarray(ver['Time'])
vert = vert/np.timedelta64(1, 's')
vers = vers/3.6

ax1 = plt.subplot(1,3,(1,2))
ax2 = plt.subplot(1,3,3)

ax1.plot(vert,vers,linewidth=0.8)
ax1.grid(which='both', alpha=0.3, linestyle='--')
ax1.set(ylabel='speed [m/s]')
ax1.set(xlabel='time [s]')

x = np.asarray(ver['X'])
y = np.asarray(ver['Y'])

k = 0
l = 0
for m in vert:
	if (m > k*5 and m < (k+1)*5):
	    ax2.plot(x[l],y[l],'o',c='#1f77b4',markersize=2)
	    ax2.annotate(str(round(vert[l],1)),(x[l]+100,y[l]-300),fontsize=5)
	    k = k+1
	    l = l+1
	else:
	    l = l+1
print(k)
print(l)
ax2.set_xticks([])
ax2.set_yticks([])
ax2.plot(x,y,linewidth=0.8)

plt.tight_layout()
plt.savefig('lap37vertrack.png',dpi=300)