import xarray as xr
import lombricquiver
import pandas as pd
# Check Version
print(lombricquiver.__version__)

ds = xr.open_dataset(
    f"https://edh:{key}@data.earthdatahub.destine.eu/era5/reanalysis-era5-single-levels-v0.zarr",
    chunks={},
    engine="zarr",
)
print(ds)

## 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_10am)

# 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(
    animation_duration=5,  # 5 seconds of total animation
    max_anchors_per_line=100,  # max number of anchors per line in the animation
    flow_speed=1.5,  # speed of the streamline agent in the animation
    virtual_time=4,  # Higher values result in longer stream lines
)
## Returns the scene class to be used in Manim
Animation_singlecolor = vf.get_scene_class()

# # Create an instance of the class animator
vf = VectorFieldAnimation()

# # Set the wind dataset in the class
vf.set_dataset(dataset_subsampled_10am)

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

## Config the animations parameters
vf.configure_streamplot(
    animation_duration = 5, # 4 seconds of total animation
    max_anchors_per_line=100 , # max number of anchors per line in the animation
    flow_speed=1.8, # it is the speed of the streamline agent in the animation,
    virtual_time=4,#Higher values therefore result in longer stream lines
)

## Config title
vf.configure_title(
    show_title = True,
    title = 'Vector Field - Wind Speed',
    font_size= 36
)

## Config suptitle
vf.configure_subtitle(
    show_subtitle = True,
    font_size=16,
    subtitle = pd.to_datetime(str(vf.dataset()['valid_time'].values)).strftime('%Y-%m-%d %H:%M:%S')
)

## Config colorbar
vf.configure_colorbar(
    show=True,
    palette = 'viridis', ## Palette - from matplotlib palettes
    num_colors = 4 ,
    figsize = (1.0,9)
)

## Check if the animator is ready 
if vf.is_ready():
    print('Everything is set')
else:
    vf.get_missing_requirements()

## Create the scene class inherented from Manim Scene
Animation = vf.get_scene_class()

%manim -ql -v WARNING Animation

## Create new background image for the animation 
new_background = create_layer_base(
    dict_extract_var,
    var_heatmap = 't2m',
    file_name = 'temp_base_layer.png'
)

## Also, let's change the time of the  day, lets plot at 17:00 
dict_extract_var = processor.extract_components_by_given_timestep(extract_variables = extract_var,
                                                                timestep=17, ## 10:00am,
                                                                lat_long=True)
dataset_subsampled_17 = dataset_subsampled.isel(valid_time=17) #17:00

vf.set_dataset(dataset_subsampled_17)
## Set the new background on the class animator
vf.set_background_image(new_background)

## Apply the same steps above

## Config the animations parameters
vf.configure_streamplot(
    animation_duration = 4, # 4 seconds of total animation
    max_anchors_per_line=75 , # max number of anchors per line in the animation
    flow_speed=1.8, # it is the speed of the streamline agent in the animation,
    virtual_time=4,#Higher values therefore result in longer stream lines
)

## Config title
vf.configure_title(
    show_title = True,
    title = 'Vector Field - Wind Speed',
    font_size= 36
)

## Config suptitle
vf.configure_subtitle(
    show_subtitle = True,
    font_size=16,
    subtitle = pd.to_datetime(str(vf.dataset()['valid_time'].values)).strftime('%Y-%m-%d %H:%M:%S')
)

## Config colorbar
vf.configure_colorbar(
    show=True,
    palette = 'viridis', ## Palette - from matplotlib palettes
    num_colors = 4 
)

## Check if the animator is ready 
if vf.is_ready():
    print('Everything is set')
else:
    vf.get_missing_requirements()

## Create the scene class inherented from Manim Scene
Animation_dif = vf.get_scene_class()

vf.dataset()

Animation.preview_config()