"""
Example script for accessing the time-resolved snapshot data and plotting the streamwise velocity.

This script demonstrates how to:
- Connect to the S3 bucket containing the dataset
- Access the Zarr-formatted data using xarray
- Extract a 2D slice of streamwise velocity data
- Visualise the data using matplotlib

Requirements:
- matplotlib
- xarray
- s3fs
- zarr
- dask (optional, for lazy loading)
"""

import matplotlib.pyplot as plt
import xarray as xr

# Set Zarr store path and S3 options
store = "s3://eidf198-highres-snapshots-sublayer-dns-tbl-re2400/data.zarr"
storage_options = {
    "anon": True,                               # Anonymous access (no credentials required)
    "endpoint_url": "https://s3.eidf.ac.uk",    # EIDF S3 endpoint
}

# Open the Zarr store snapshots group as an xarray dataset
ds = xr.open_zarr(
    store,
    consolidated=True,  # Use consolidated metadata for faster access
    storage_options=storage_options,
    group="snapshots",  # Access the snapshots group in the Zarr store
)

# Extract streamwise velocity from the first snapshot in the x-z plane
# If using Dask, this will be lazy-loaded until explicitly computed
x_slice = slice(0, 456)                 # Select a subset of the x-dimension for easier visualisation
x = ds["x"].isel(x=x_slice)             # Streamwise coordinate
z = ds["z"]                             # Spanwise coordinate
u = ds["u"].isel(t=0, x=x_slice, y=-1)  # Streamwise velocity at the last y-position

# Plot (matplotlib implicitly calls compute() on the Dask array when plotting)
fig, ax = plt.subplots(figsize=(8, 3))
contour = ax.contourf(x, z, u.T, levels=100)

# Add colorbar and labels
cbar = fig.colorbar(contour, ax=ax)
cbar.set_label("Streamwise Velocity (u)")
ax.set_xlabel("Streamwise Coordinate (x)")
ax.set_ylabel("Spanwise Coordinate (z)")
ax.set_title("Streamwise Velocity in X-Z Plane")
ax.set_aspect("equal")
fig.tight_layout()
plt.show()