"""
Example script for accessing the time-averaged statistics data and converting to inner (viscous) units.

This script demonstrates how to:
- Connect to the S3 bucket containing the dataset
- Access the Zarr-formatted data using xarray
- Extract root level metadata (e.g. kinematic viscosity)
- Extract a subset of the streamwise velocity data
- Convert velocity and coordinate data to inner (viscous) units
- 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 root group as an xarray dataset and extract viscosity
ds = xr.open_zarr(
    store,
    consolidated=True,  # Use consolidated metadata for faster access
    storage_options=storage_options,
)
viscosity = ds.attrs["viscosity"]

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

# Extract wall-normal coordinate and streamwise velocity at a specific streamwise location
# If using Dask, this will be lazy-loaded until explicitly computed
y = ds["y"]                 # Wall-normal coordinate
u = ds["u"].sel(x=500.0)    # Streamwise velocity at a specific streamwise location

# Average over the (homogeneous) spanwise coordinate
u = u.mean(dim="z")

# Compute friction velocity (u_infty = 1 and rho_infty = 1)
tau = 1.0 * viscosity * (u.isel(y=1) - u.isel(y=0)) / (y[1] - y[0]) # Wall shear stress
utau = (tau / 1.0) ** 0.5                                           # Friction velocity

# Convert velocity and coordinate data to inner (viscous) units
y_plus = y / (viscosity / utau) # Wall-normal coordinate in inner units
u_plus = u / utau               # Streamwise velocity in inner units

# Plot (matplotlib implicitly calls compute() on the Dask array when plotting)
fig, ax = plt.subplots(figsize=(8, 5))
ax.plot(y_plus, u_plus)

# Set axis limits and scale
ax.set_xlim(0.5, 4000)
ax.set_xscale("log")

# Add labels
ax.set_xlabel("Wall-Normal Coordinate (y+)")
ax.set_ylabel("Streamwise Velocity (u+)")
ax.set_title("Time/Spanwise-Averaged Streamwise Velocity along Wall-Normal Coordinate")
fig.tight_layout()
plt.show()