#!/usr/bin/env python
# coding: utf-8

# ## Install/Update packages and Load common functions

# In[ ]:


# only run once, then restart session if needed
get_ipython().system('pip install uv')

import os
import sys

def is_colab():
    try:
        import google.colab
        return True
    except ImportError:
        return False

if is_colab():
    os.system('uv pip install --system -r https://raw.githubusercontent.com/aodn/aodn_cloud_optimised/main/notebooks/requirements.txt')
else:
    os.system('uv venv')
    os.system('uv pip install -r https://raw.githubusercontent.com/aodn/aodn_cloud_optimised/main/notebooks/requirements.txt')


# In[1]:


import xarray as xr
import fsspec


# In[2]:


# remote zarr dataset
dataset_name="radar_TurquoiseCoast_velocity_hourly_average_delayed_qc"
url = f's3://imos-data-lab-optimised/cloud_optimised/cluster_testing/{dataset_name}.zarr/'
ds = xr.open_zarr(fsspec.get_mapper(url, anon=True), consolidated=True)
ds


# In[3]:


from matplotlib.pyplot import figure, pcolor, colorbar, xlabel, ylabel, title, draw, quiver, show
import matplotlib.pyplot as plt
import numpy as np

def acorn_plot(ds, time_start):
    """
    Plotting function for ACORN data.

    Parameters:
    - ds: xarray dataset
        The input dataset.
    - time_start: str
        The starting time in the format '2021-02-21T01:00:00'.
    """
    # Create a 3x2 grid of subplots
    fig, axes = plt.subplots(nrows=3, ncols=2, figsize=(10, 10))

    # Plot data on each subplot. We're plotting the next 6 time stamps (requiring obviously that they all exist)
    ii = 0
    iTime = list(ds.TIME.values).index(ds.sel(TIME=time_start , method='nearest').TIME)
    cbar_ax = fig.add_axes([0.99, 0.1, 0.02, 0.8])  # Adjust the position and size of the colorbar

    for i in range(3):
        for j in range(2):
            uData = ds.UCUR[iTime + ii, :, :]
            vData = ds.VCUR[iTime + ii, :, :]
            speed = np.sqrt(uData**2 + vData**2)
            lonData = ds.LONGITUDE.values
            latData = ds.LATITUDE.values

            p = axes[i, j].pcolor(lonData, latData, speed)
            axes[i, j].quiver(lonData, latData, uData, vData, units='width')
            axes[i, j].set_title(f'{np.datetime_as_string(ds.TIME.values[iTime + ii])}')
            ii += 1

    # Add a common colorbar
    fig.colorbar(p, cax=cbar_ax, label='Speed')

    # Adjust layout for better appearance
    plt.tight_layout()

    # Show the plot
    plt.show()


# In[5]:


get_ipython().run_cell_magic('time', '', "acorn_plot(ds, '2020-01-01T00:00:00')\n")


# In[ ]:





# In[ ]: