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

# In[1]:


import xarray as xr
import earthaccess
from dask.distributed import Client
import earthaccess.virtualizarr as evz


# In[2]:


client = Client()
client


# # Opening multiple granules with direct cloud access

# In[3]:


get_ipython().run_cell_magic('time', '', 'results = earthaccess.search_data(\n    count=100,\n    short_name="MUR-JPL-L4-GLOB-v4.1"\n)\n')


# In[4]:


get_ipython().run_cell_magic('time', '', 'vds = evz.open_virtual_mfdataset(results, access="direct", concat_dim="time", coords="all", compat="override", combine_attrs="drop_conflicts")\nvds\n')


# In[5]:


s3fs = earthaccess.get_s3fs_session(results=results)
options = s3fs.storage_options.copy()


# In[6]:


get_ipython().run_cell_magic('time', '', 'refs = vds.virtualize.to_kerchunk(filepath=None, format="dict")\noptions["fo"] = refs\nmur = xr.open_dataset(\n    "reference://",\n    engine="zarr",\n    chunks={},\n    backend_kwargs={"storage_options": options, "consolidated": False},\n)\nmur\n')


# In[7]:


print(f'{mur.nbytes / 1e12} TB')


# In[8]:


get_ipython().run_cell_magic('time', '', 'mur.isel(time=0).sel(lat=slice(15,40), lon=slice(-95, -50)).analysed_sst.plot.pcolormesh(x="lon", y="lat", cmap="viridis")\n')


# # Save virtual reference file and load with xarray

# In[10]:


filepath = "MUR-JPL-L4-GLOB-v4.1.json"
vds.virtualize.to_kerchunk(filepath=filepath, format="json")


# In[11]:


options["fo"] = filepath
mur = xr.open_dataset(
    "reference://",
    engine="zarr",
    chunks={},
    backend_kwargs={"storage_options": options, "consolidated": False},
)
mur


# # Open a single file dataset

# In[12]:


results = earthaccess.search_data(
    count=5,
    short_name="SWOT_L2_LR_SSH_Expert_2.0"
)


# In[13]:


get_ipython().run_cell_magic('time', '', 'vds = evz.open_virtual_dataset(results[0], access="direct")\nvds\n')


# In[14]:


s3fs = earthaccess.get_s3fs_session(results=results)
options = s3fs.storage_options.copy()


# In[15]:


refs = vds.virtualize.to_kerchunk(filepath=None, format="dict")
options["fo"] = refs
swot = xr.open_dataset(
    "reference://",
    engine="zarr",
    chunks={},
    backend_kwargs={"storage_options": options, "consolidated": False},
)
swot


# # Create virtual references without direct cloud access (use indirect/https)

# In[18]:


get_ipython().run_cell_magic('time', '', 'results = earthaccess.search_data(\n    count=100,\n    short_name="AVHRR_OI-NCEI-L4-GLOB-v2.1"\n)\n')


# In[19]:


get_ipython().run_cell_magic('time', '', 'vds = evz.open_virtual_mfdataset(results, access="indirect", concat_dim="time", coords="all", compat="override", combine_attrs="drop_conflicts")\nvds\n')


# In[20]:


filepath = "AVHRR_OI-NCEI-L4-GLOB-v2.1.json"
vds.virtualize.to_kerchunk(filepath=filepath, format="json")