import cartopy.crs as ccrs
import matplotlib.pyplot as plt
from pyproj import CRS
import pygrib
import cfgrib
import xarray as xr
import sys
import metpy

# First, you should download a HRDPS grib2 file from
# https://dd.weather.gc.ca/model_hrdps/continental/2.5km/00/00/
# I grabbed the TMP_AGL-2m file...
FILE = "20230722T00Z_MSC_HRDPS_TMP_AGL-2m_RLatLon0.0225_PT000H.grib2"
ds = xr.open_dataset(FILE)

# Parse the projparams with pygrib, because cfgrib doesn't do this for you
# (see https://github.com/ecmwf/cfgrib/issues/251)
with pygrib.open(str(FILE)) as grb:
    msg = grb.message(1)
    cf_params = CRS(msg.projparams).to_cf()

ds["gribfile_projection"] = None
ds["gribfile_projection"].attrs = cf_params
for var in list(ds):
    if var == "gribfile_projection":
        continue
    ds[var].attrs["grid_mapping"] = "gribfile_projection"

# Ok, now we are ready to check that the changes to metpy/mapping.py worked...
variables = [i for i in list(ds) if len(ds[i].dims) > 0]
ds = ds.metpy.parse_cf(varname=variables)
crs = ds.metpy_crs.item().to_cartopy()

# And plot the data...
ax = plt.subplot(111, projection=crs)
ax.coastlines()
ax.pcolormesh(ds.longitude, ds.latitude, ds.t2m, transform=ccrs.PlateCarree())
ax.set_global()
ax.gridlines()