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

# # ReferenceMaker/ReferenceFileSystem GRIB2/HRRR Example
# 
# Requires development version of `fsspec_reference_maker`
# * `pip install --user git+https://github.com/intake/fsspec-reference-maker`

# In[1]:


import json
import fsspec
from fsspec_reference_maker.grib2 import scan_grib
import os


# ## Create reference jsons

# In[ ]:


# 1GB of data files, forming a time-series
filter={'typeOfLevel': 'heightAboveGround', 'level': 2}

files = ['s3://noaa-hrrr-bdp-pds/hrrr.20190101/conus/hrrr.t22z.wrfsfcf01.grib2',
         's3://noaa-hrrr-bdp-pds/hrrr.20190101/conus/hrrr.t23z.wrfsfcf01.grib2',
         's3://noaa-hrrr-bdp-pds/hrrr.20190102/conus/hrrr.t00z.wrfsfcf01.grib2',
         's3://noaa-hrrr-bdp-pds/hrrr.20190102/conus/hrrr.t01z.wrfsfcf01.grib2',
         's3://noaa-hrrr-bdp-pds/hrrr.20190102/conus/hrrr.t02z.wrfsfcf01.grib2',
         's3://noaa-hrrr-bdp-pds/hrrr.20190102/conus/hrrr.t03z.wrfsfcf01.grib2',
         's3://noaa-hrrr-bdp-pds/hrrr.20190102/conus/hrrr.t04z.wrfsfcf01.grib2',
         's3://noaa-hrrr-bdp-pds/hrrr.20190102/conus/hrrr.t05z.wrfsfcf01.grib2',
         's3://noaa-hrrr-bdp-pds/hrrr.20190102/conus/hrrr.t06z.wrfsfcf01.grib2']
so = {"anon": True, "default_cache_type": "readahead"}
common = ['time', 'step', 'latitude', 'longitude', 'valid_time']


# In[ ]:


def create_jsons(files):
    for url in files:
        out = scan_grib(url, common, so, inline_threashold=100, filter=filter)
        with open(os.path.basename(url).replace("grib2", "json"), "w") as f:
            json.dump(out, f)


# In[6]:


create_jsons(files)


# ## Use `MultiZarrToZarr()` to combine into single reference

# In[7]:


from glob import glob
json_list = sorted(glob('./hrrr.t*.json'))


# In[10]:


from fsspec_reference_maker.combine import MultiZarrToZarr

mzz = MultiZarrToZarr(json_list, remote_protocol="s3", remote_options={"anon": True},
                        xarray_concat_args={"dim": 'time'})
mzz.translate("hrrr.total.json")


# ## Access data and plot

# In[23]:


import xarray as xr
from fsspec_reference_maker.grib2 import GRIBCodec
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.pyplot as plt
import pandas as pd


# In[3]:


fs = fsspec.filesystem('reference', fo='./hrrr.total.json', remote_protocol='s3', remote_options={'anon':True})
m = fs.get_mapper('')
ds = xr.open_dataset(m, engine='zarr')


# In[4]:


ds


# In[25]:


i = 0 # Time index

fig = plt.figure(figsize=(6,6))
ax = plt.subplot(111, projection=ccrs.LambertConformal())
p = ax.pcolormesh(ds.longitude, ds.latitude, ds.si10.isel(time=i), transform=ccrs.PlateCarree())
ax.coastlines()
ax.add_feature(cfeature.STATES)

time = pd.to_datetime(ds.time[i].data).strftime("%Y-%m-%d %H%M UTC")

ax.set_title(f"10m Wind Speed\n{time}")
plt.colorbar(p, orientation='horizontal', label='m/s')


# In[ ]: