"""
Example script for accessing the time-averaged statistics data and plotting the skin friction profile.

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)
- Compute the skin friction coefficient along the streamwise direction
- 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 coordinates and streamwise velocity
# If using Dask, this will be lazy-loaded until explicitly computed
x = ds["x"] # Streamwise coordinate
y = ds["y"] # Wall-normal coordinate
u = ds["u"] # Streamwise velocity

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

# Compute skin friction coefficient (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
cf = tau / (0.5 * 1.0 * 1.0**2)

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

# Set axis limits
ax.set_xlim(40.0, 950.0)

# Add labels
ax.set_xlabel("Streamwise Coordinate (x)")
ax.set_ylabel("Skin Friction Coefficient (Cf)")
ax.set_title("Skin Friction Coefficient along Streamwise Coordinate")
fig.tight_layout()
plt.show()