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

# # HRRR forecast best time series
# Read from kerchunked GRIB2 files, write subset to Zarr dataset

# In[1]:


import xarray as xr
import fsspec
import hvplot.xarray
from kerchunk.grib2 import scan_grib   # needed here only for grib compression codec


# #### Open HRRR referenceFileSystem dataset

# In[2]:


rpath = 's3://esip-qhub-public/noaa/hrrr/hrrr_best.json'
s_opts = {'anon':True, 'skip_instance_cache':True}
r_opts = {'anon':True}
fs = fsspec.filesystem("reference", fo=rpath, ref_storage_args=s_opts,
                       remote_protocol='s3', remote_options=r_opts)
m = fs.get_mapper("")

ds = xr.open_dataset(m, engine="zarr", backend_kwargs=dict(consolidated=False), 
                      chunks={'valid_time':1})


# In[3]:


ds = ds.drop(['time', 'step', 'heightAboveGround']).rename({'valid_time':'time'})


# In[4]:


ds


# In[5]:


da = ds.isel(time=slice(-3,-1))['t2m'].load()
da.hvplot.quadmesh(x='longitude', y='latitude', cmap='turbo',
                        geo=True, tiles='OSM', rasterize=True)


# #### Write subset dataset to Zarr

# In[6]:


ds.isel(time=slice(-3,-1)).to_zarr('foo.zarr', 'w')