start_time = time.perf_counter()   

## Create fig
fig, ax = plt.subplots(figsize=(12.2, 6.8),
                       dpi=100,
                       subplot_kw={'projection': ccrs.PlateCarree()})

## Add features to the map 
ax.add_feature(cfeature.COASTLINE, linewidth=0.5)
ax.add_feature(cfeature.BORDERS, linewidth=0.5)
ax.add_feature(cfeature.LAND)

## Labels 
plt.xlabel('Longitude')
plt.ylabel('Latitude')

## Create a heatmap based on the velocity 
## Use the whole data and not the subsampled one.
heatmap = ax.pcolormesh(dict_extract_var['long'],
                        dict_extract_var['lat'],
                        dict_extract_var['wind_speed'],
                        cmap="plasma_r",
                        transform=ccrs.PlateCarree())

# Initialize the streaamplot
stream = ax.streamplot(dataset_subsampled['longitude'].values,
                       dataset_subsampled['latitude'].values,
                       dataset_subsampled['u10'],
                       dataset_subsampled['v10'],
                       color="white", density=10, linewidth=0.4)


#Colorbar
pos = ax.get_position()
cbar_ax = fig.add_axes([pos.x1 + 0.01, pos.y0, 0.02, pos.height]) #axes of colorbar 
cbar = fig.colorbar(heatmap, cax=cbar_ax, orientation='vertical')

#Labels and Title 
cbar.set_label("[m/s]")
plt.title('ERA5 - V10')

# get the geographic bounds:
fig_lon_min, fig_lon_max, fig_lat_min, fig_lat_max = ax.get_extent(crs=ccrs.PlateCarree())

lengths = []
colors = []
lines = []

## Get segments to loop through it 
segments = stream.lines.get_segments()

for streamline in segments:
    # Transpose and retrieve lat and long
    s = streamline.T
    x, y = s[0], s[1]

    # Create a vector of streamline points 
    points = np.array([x, y]).T.reshape(-1, 1, 2)
    
    # Concat along second dimension (column)
    ## Basically is adding a second column of a shift vector 
    seg = np.concatenate([points[:-1], points[1:]], axis=1)

    n = len(seg)
    # Compute cumulative length along streamline
    D = np.sqrt(((points[1:] - points[:-1])**2).sum(axis=-1)) #compute length 
    L = D.cumsum().reshape(n, 1) + np.random.uniform(0, 1) #cumulative sum 

    # Create white color with gradient transparency
    C = np.ones((n, 4))  # (R=1, G=1, B=1, Alpha)
    C[:, 3] = np.linspace(0.0, 1.0, n)
    C[::-1] = (L*1.5) % 1

    # Apply colors and transparency to LineCollection
    line = LineCollection(seg, colors=C, linewidth=0.4, transform=ccrs.PlateCarree())
    lengths.append(L)
    colors.append(C)
    lines.append(line)
    ax.add_collection(line)

def update(frame_no):
    for i in range(len(lines)):
        # Adjust movement based on random wind speed (scaled appropriately)
        lengths[i] -= random.uniform(0.1, 1.2)
        # Adjust transparency dynamically
        colors[i][::-1] = (lengths[i] * 0.1) % 1
        # Update streamline color
        lines[i].set_color(colors[i][::-1])
    pbar.update()

n = 50
animation = FuncAnimation(fig, update, frames=n, interval=8)
pbar = tqdm.tqdm(total=n)
# Using FFmpeg writer
writer = FFMpegWriter(fps=8, bitrate=1800)
animation.save('Matplotlib_era5_wind.mp4', writer=writer)
pbar.close()
plt.close()

end_time = time.perf_counter()
print(f"Time taken to create the plot: {end_time - start_time:.2f} seconds")

## Visualize
import mediapy as media
video = media.read_video('Matplotlib_era5_wind.mp4')
media.show_video(video, title = "Wind Vector Fields in Europe", fps=8, width=700)

start_time = time.perf_counter()

## Create the Animator instance
vf = VectorFieldAnimation()

## All the following code are methods of the VectorFieldAnimation class
# Set the wind dataset in the class
vf.set_dataset(dataset_subsampled)

# Create the background image and set under the hood
vf.create_background_image(
    dict_extract_var=dict_extract_var,
    var_heatmap = 'wind_speed' 
)

vf.configure_streamplot(
    delta_y=0.15,
    resize_factor=0.05,
    stroke_width=1.0,
    flow_speed=1.5,
    time_width=0.3,
    animation_duration=5,
    dt=0.01,
    max_anchors_per_line=150,
    color="WHITE",
    virtual_time=4
)
## Returns the scene class to be used in Manim
Animation = vf.get_scene_class()

%manim -ql -v WARNING Animation

end_time = time.perf_counter()
print(f"Time taken to create the plot: {end_time - start_time:.2f} seconds")