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

# # **grib2io v2.0 Demo**
# 
# This notebook demos some of the new code design and features of grib2io v2.0.  The sample file used is a GRIB2 files containing GFS model forecasts from the 00Z cycle on a 0.25 deg. global lat/lon grid with a lead time of 6 hours. A similar sample file can be downloaded from [here](http://nomads.ncep.noaa.gov/pub/data/nccf/com/gfs/prod/)

# In[1]:


import grib2io
import numpy as np


# ## **Opening GRIB2 files**
# 
# Note that nothing has changed between v1 and v2 with respect to to opening GRIB2 files and the same attributes exist.

# In[2]:


g = grib2io.open("gfs.t00z.pgrb2.0p25.f006")
print(g)


# ### Here are some key things to know about what grib2io is doing when opening a GRIB2 file.
# * The messages in the file are indexed automatically.  In other words, GRIB2 sections 0, 1, 2, 3, 4, and 5 are read, unpacked, and used to create Grib2Message objects.
# * Message, bitmap, and data offsets are recorded.
# * GRIB2 messages that contain submessages are automatically flattened.
# * At this point, bitmap and data sections are not read.  These sections make up the majority of packed message, so we defer reading these parts of the message until absolutely necessary.

# ## **Investigation of the file object index**
# 
# Underneath the hood, the indexing has changed significantly. There no longer is GRIB2-specific metadata from messages, but just information about location
# and size of messages.  We are now storing the GRIB2 message object in the index in key 'msg'.
# 
# One might think that storing all of these objects in the index would negatively impact performance, but the new GRIB2 message object is extremely lightweight (more on that later) and acts like a templated container of metadata.  At this point, the data section is not a part of the GRIB2 message object.

# In[3]:


print(g._index.keys())


# ## **Reading GRIB2 Messages from file**

# ### Use read(), seek(), tell()
# 
# The `grib2io.open` class provides custom implementations of read(), seek(), and tell().  While these methods are traditionally used for traversing bytes, these custom implementations operate on units of GRIB2 messages.  Note that when these methods are used, you are not acutally performning any I/O to the GRIB2 file.  You are positioning inside the list contained by `grib2io.open._index['msg']`.  The following example puts the "file object" pointer to the 100th message and reads the next 15 messages (from message 100).  In this example, there is no file IO -- just communicating with the Grib2Message objects stored in the index.

# In[4]:


g.seek(100)
print('GRIB2 file poistion: ',g.tell())
msgs = g.read(15)
print('Number of messages read: ',len(msgs))
for msg in msgs:
    print(msg)


# ### Use index, slice, or key notation on the file object
# 
# **_IMPORTANT: In grib2io v2.0, you no longer have to specify a 2nd index when select 1 message !!_**

# In[5]:


# Index
msg = g[50]
print(msg)


# **_IMPORTANT:  In grib2io v2.0, a 0th message now exists !!_**

# In[3]:


# Index
msg = g[0]
print(msg)


# In[7]:


# Slice
msgs = g[80:90]
for msg in msgs:
    print(msg)


# In[8]:


# Key - where string is the GRIB2 shortName attribute
msgs = g['CAPE']
for msg in msgs:
    print(msg)


# ### Selection of GRIB2 messages based on attributes
# 
# Use **_any_** Grib2Message object attribute name as a keyword for the file object select() method. **NOTE:** The select() method will always return a list.

# #### Example 1: Select temperature message (TMP) that use product definition template number of 8.

# In[9]:


msgs = g.select(shortName='TMP',productDefinitionTemplateNumber=8)
for msg in msgs:
    print(msg)
    print('\t',msg.duration,msg.statisticalProcess)


# #### Example 2: Select U-wind messages (UGRD) at geometric heights above ground.

# In[10]:


msgs = g.select(shortName='UGRD',typeOfFirstFixedSurface=103)
for msg in msgs:
    print(msg)
    print('\t',msg.typeOfFirstFixedSurface)


# ## **Grib2Message Object Metadata Attributes**
# 
# The Grib2Message object, attributes, and metadata storage have changed **significantly** from v1.0.0 or v2.0.0.  The Grib2Message class inherits from a Grib2Message base class that contains attributes for sections that are not templates.  GRIB2 Sections 3, 4, and 5 are defined by templates and therefore the attributes can be different between templates of the same section.  All template classes are defined as a `dataclasses.dataclass` and each attribute within each template class is a class variable and defined to be a `dataclasses.field` object.  Further, each attribute, defined as a field has a default value that is a unique descriptor class with custom implementations of the `__get__` and `__set__` protocols.
# 
# Below is an example of all of the metadata associated with a GRIB2 message for a GFS 500mb HGT field.  Each metadata item you see is being created dynamically from the numeric GRIB2 section metadata.  For example, in section 4, there is the attribute `parameterCategory`.  This attribute has a custom descriptor class, `grib2io.templates.ParameterCategory`, whose `__get__` protocol points to the appropriate index location in the section 4 array.  If we wanted to change this attribute value, you can perform `msg.parameterCategory = 5`.  This would trigger the `__set__` protocol for the `grib2io.templates.ParameterCategory` descriptor class which would put that new value into the appropriate index location in the section 4 array.

# In[4]:


msg = g.select(shortName='HGT',level='500 mb')[0]
msg


# ### Grib2Message wgrib2-like `__str__`
# 
# The Grib2Message object provides a wgrib2-like 1-liner string for `__str__` (as shown above).  The Grib2Message object `__repr__` will provide all metadata (see above).

# In[5]:


print(msg)


# ### Grib2Message Construction

# In[13]:


print('Section 0: ',msg.section0)
print('Section 1: ',msg.section1)
print('Section 2: ',msg.section2)
print('Section 3: ',msg.section3)
print('Section 4: ',msg.section4)
print('Section 5: ',msg.section5)


# The Grib2Message class is dynamically created and is unique to GRIB2's grid definition, production definition, and data representation template numbers. These templates contain metadata attributes unique to the definition of the given template.
# 
# Here we can use the class method resolution order (`__mro__`) to show the various classes that make up the Grib2Message class. 

# In[14]:


for c in msg.__class__.__mro__:
    print(c)


# The template classes listed above contain the attributes as class variables that are unique to its template.
# 
# You will notice that the `ProductDefinitionTemplate0` is blank.  This is because it contains the common attributes used by the rest of the PTDNs, so they are defined in class `_Grib2Message`.

# In[15]:


print(grib2io.templates.GridDefinitionTemplate0._attrs)
print(grib2io.templates.ProductDefinitionTemplate0._attrs)
print(grib2io.templates.DataRepresentationTemplate3._attrs)


# Now lets take a look at how the metadata attributes are linked to their origin locations in a given section.  Here we will look at the Grib2Message `refDate` attribute, whose components are stored in Section 1.
# 
# To reiterate, `msg.refDate` does not hold a value that would be a `datetime.datetime` object.  The `refDate` attribute is defined as a descriptor class, specifically `grib2io.templates.RefDate()`, whose `__get__` protocol returns a `datetime.datetime` object generated from the date components in section 1.

# In[16]:


print(msg.section1)
print(msg.refDate)
print(msg.year,msg.month,msg.day,msg.hour,msg.minute,msg.second)


# To demonstrate that these attributes are dynamically created, lets change the year and minute. We do this by setting the attributes to new values.  You will see that values have changed in section1 and other attributes.

# In[17]:


msg.year = 2023
msg.minute = 55
print(msg.section1)
print(msg.refDate)


# **IMPORTANT:** In grib2io v2.0, all date and time attributes are `datetime.datetime` or `datetime.timedelta` objects.  Let use precipitation as an example.

# In[18]:


msg = g['APCP'][0]
print(msg)
print(msg.refDate,type(msg.refDate))
print(msg.validDate,type(msg.validDate))
print(msg.leadTime,type(msg.leadTime))
print(msg.duration,type(msg.duration))


# This allows for datetime arithmetic
# 
# The Grib2Message object defines the property `pdy` that provides the reference date as a string in `'YYYYMMDD'` format.

# In[19]:


print(msg.refDate + msg.leadTime)
print(msg.validDate)
print(msg.refDate+msg.leadTime==msg.validDate)
print(msg.pdy)


# The Grib2Message object provides a method, `attrs_by_section()` to get attribute names and values (if desired) for a given GRIB2 section.

# In[20]:


for k,v in msg.attrs_by_section(3,values=True).items():
    print(k,v)


# Any GRIB2 metadata attribute that is a code value means that it is associated with a code table. That attribute is likely to be of type `grib2io.templates.Grib2Metadata` which is a class to hold the code value, the definition given its lookup table.
# 
# Let use the GRIB2 metadata attribute `discipline` as an example.

# In[21]:


print(type(msg.discipline))
print(msg.discipline)
print(msg.discipline.value)
print(msg.discipline.definition)
print(msg.discipline.table)


# ## A Quick Detour to NCEP GRIB2 Tables
# 
# grib2io contains the [NCEP GRIB2 Tables](https://www.nco.ncep.noaa.gov/pmb/docs/grib2/grib2_doc/) stored as python dictionaries in `grib2io.tables`.  You are able to access them directly,

# In[22]:


print(grib2io.tables.section0.table_0_0)
print(grib2io.tables.section4.table_4_6)


# However, it is more useful to use the various helper functions in `grib2io.tables` to view the NCEP GRIB2 code tables and lookup values.

# In[23]:


for k,v in grib2io.tables.get_table('0.0').items():
    print(f'{k}: {v}')


# In[24]:


print('Value 2 from table 4.6 means',grib2io.tables.get_value_from_table(2,'4.6'))


# In[25]:


discipline, parmcat, parmnum = 0, 1, 10
varinfo = grib2io.tables.get_varinfo_from_table(discipline,parmcat,parmnum)
print(f'Discipline value of {discipline} with paramater category and number of {parmcat} and {parmnum} defines variable {varinfo}')


# Some metadata attributes exist for more easier understanding of information from the Grib2Message.
# 
# The `level` attribute will provide a layer/level string that is **_exactly the same as wgrib2_** and is generated by `grib2io.tables.get_wgrib2_level_string()`.

# In[26]:


print(msg.level)
print(msg.units) # Units of the variable
print(msg.unitOfFirstFixedSurface) # Units of the level (or layer)


# ## **Unpacking Data from Grib2Message Objects**
# 
# Lets go back to the 500 mb HGT example.
# 
# Data are accessible via the `data` attribute.
# 
# **IMPORTANT:** If you unfamiliar with Python/NumPy, by default, multi-dimensional array data are ordered in row-major order.  This is how C/C++ order array data in memory.  Fortran orders array data in column-major order. So 2D array will be ordered (ny, nx).

# In[27]:


msg = g.select(shortName='HGT',level='500 mb')[0]
print(msg)
hgtdata = msg.data
print(hgtdata)
print(hgtdata.shape)


# ## **Geospatial Information from Grib2Message Objects**
# 
# GRIB2 Section 3 contains all of the geospatial information about the message.

# In[28]:


print(msg.section3)
for k,v in msg.attrs_by_section(3,values=True).items():
    print(f'{k}: {v}')


# grib2io also provides the coorindate system in terms of PROJ parameters.

# In[29]:


print(msg.projParameters)


# ### Generating Latitude and Longitude Value
# 
# Latitude and Longitude values can be generated via `grid()` or `latlons()` methods.

# In[30]:


lats, lons = msg.grid()
print(lats)
print(lons)
msg.grid() == msg.latlons() # latlons() is an alias of grid()


# **NEW** In grib2io v2.0, computed lats and lons are stored so they do not have to regenerated if you call the methods from another message with the same coordinate system.  This is determined by taking the SHA-1 of section 3 and storing that as a key in a dictionary where lats and lons are the values.

# In[31]:


print(msg._sha1_section3)
print(grib2io._grib2io._latlon_datastore)


# ## **Interpolation**
# 
# **NEW** In grib2io v2.0 is the ability to perform spatial interpolation from one grid to another by interfacing to the [NCEPLIBS-ip](https://github.com/NOAA-EMC/NCEPLIBS-ip) library.
# 
# The Grib2Message object already contains the necessary information about the input grid. So all we need to go here is to define the output grid (i.e. the grid that we want to interpolate to).  We do that by specifcy the GRIB2 Grid Definition Template Number and Template values.  Here we will interpolate to the NBM CONUS 2.5km Lambert Conformal grid.
# 
# Here we place the template number and template into a `Grib2GridDef` object.

# In[32]:


nbm_co_gdtn = 30
nbm_co_gdt = [        1,         0,   6371200,       255,       255,       255,
                    255,      2345,      1597,  19229000, 233723400,        48,
               25000000, 265000000,   2539703,   2539703,         0,        80,
               25000000,  25000000, -90000000,         0]
grid_def_out = grib2io.Grib2GridDef(nbm_co_gdtn,nbm_co_gdt)
print(grid_def_out.gdtn)
print(grid_def_out.gdt)


# To interpolate a Grib2Message object, call the `interpolate()` method.  The return from the method is a new Grib2Message object.  All metadata remain the same except for section 3 metadata.
# 
# **NOTE 1: The interpolation option can be a string or numeric value determined by the NCEPLIBS-ip library.**
# 
# **NOTE 2: The message number of interpolated messages is -1. This is becasue this message does not originate from a physical file.  Any Grib2Message object created from scatch will have message number of -1 until it is associated with a file object.**

# In[33]:


new = msg.interpolate('bilinear',grid_def_out)
print(new)
print(new.section3)
for k,v in msg.attrs_by_section(3,values=True).items():
    print(f'{k}: {v}')
print(new.data)


# ## **Creating Grib2Message Object from Scratch**

# The majority of the time, you will most likely have an existing Grib2Message object that you could use a guideline or template in building a new Grib2Message object.  To demo, this we are going to use a bit of information from scratch and from an existing Grib2Message.
# 
# We will construct a Grib2Message object containing the **10-m Wind Speed from the U- and V-wind components**.

# In[34]:


uwind = g.select(shortName='UGRD',level='10 m above ground')[0]
vwind = g.select(shortName='VGRD',level='10 m above ground')[0]
print(uwind)
print(vwind)


# Create new Grib2Message object.
# 
# Here we generate the message by only passing Grid Definition Template Number, Product Definition Template Number, and Data Representation Template Number which we can borrow from either of the `uwind` or `vwind` objects.

# In[36]:


wspd = grib2io.Grib2Message(gdtn=uwind.gdtn,pdtn=uwind.pdtn,drtn=uwind.drtn)
print(wspd)


# Notice how printing the new object results in an error, specifically the `refDate`.  Since the GRIB2 message object was created from scratch it does have the appropriate metadata filled in yet.
# 
# Lets take at the sections 0-5.  Lets get a list of attributes and then populate them.  To perform this correctly, you need to some familiarity with the NCEP GRIB2 tables.
# 
# ### Section 0

# In[37]:


print(wspd.attrs_by_section(0))
wspd.discipline = 0 # 0 = Meteorological Products
print(wspd.attrs_by_section(0, values=True))
print(wspd.section0)


# ### Section 1
# 
# Here we are going to cheat a bit by populating section 1 attributes from the uwind variable.  The following simulates populating the attributes by hand

# In[38]:


for attr in wspd.attrs_by_section(1):
    try:
        cmd = 'wspd.'+attr+' = uwind.'+attr+'.value'
        exec(cmd)
    except(AttributeError):
        cmd = cmd.replace('.value','')
        exec(cmd)

for k,v in wspd.attrs_by_section(1,values=True).items():
    print(f'{k}: {v}')


# ### Section 3
# 
# For the rest of the sections, we will simply populate from `uwind` variable and adjust as needed.

# In[39]:


wspd.section3 = uwind.section3
for k,v in wspd.attrs_by_section(3,values=True).items():
    print(f'{k}: {v}')


# ### Section 4
# 
# Here we need to modify the parameter Category and Number to define the Wind Speed

# In[40]:


wspd.section4 = uwind.section4
wspd.parameterCategory = 2 # Momentum, https://www.nco.ncep.noaa.gov/pmb/docs/grib2/grib2_doc/grib2_table4-1.shtml
wspd.parameterNumber = 1 # Wind Speed, https://www.nco.ncep.noaa.gov/pmb/docs/grib2/grib2_doc/grib2_table4-2-0-2.shtml
for k,v in wspd.attrs_by_section(4,values=True).items():
    print(f'{k}: {v}')


# ### Section 5

# In[41]:


wspd.section5 = uwind.section5
for k,v in wspd.attrs_by_section(5,values=True).items():
    print(f'{k}: {v}')
print(wspd.section5)


# **IMPORANT:** So we set section 5 of the `wspd` variable to section 5 of `uwind` variable. This might look harmless, but notice that there is specific information in section 5 related to the packing of the U-wind data.  Having this data present in section 5 for `wspd` could lead to undesirable effects when packing.  The majority of attributes in section 5 will be set by the packing code. So the following will we will sanitze section 5.
# 
# **NOTE:** The sanitization will likely become an under-the-hood feature of grib2io.

# In[64]:


wspd.section5 = np.zeros(wspd.section5.shape,dtype=wspd.section5.dtype)
wspd.section5[0] = wspd.nx * wspd.ny
wspd.section5[1] = uwind.dataRepresentationTemplateNumber.value
wspd.typeOfValues = 0 # Float
wspd.binScaleFactor = 0 # Binary Scale Factor
wspd.decScaleFactor = 3 # Decimal Scale Factor
wspd.groupSplittingMethod = 1 # Group Split
wspd.spatialDifferenceOrder = 2 # Second order spatial differencing
print(wspd.section5)


# We have now defined the packing of the `wspd` variable. We will use the same packing scheme as the `uwind` (Complex Packing with Spatial Differencing)

# ### Packing the Data and Message
# 
# First we need to compute the wind speed from the U- and V-Wind components.

# In[46]:


wspd.data = np.sqrt(uwind.data**2.0+vwind.data**2.0)
print(np.amin(wspd.data),np.amax(wspd.data))


# Next, we can call the `pack()` method to pack the data and create the packed GRIB2 binary message.

# In[56]:


wspd.pack()


# We now have a packed GRIB2 message.  It is stored in the Grib2Message object in the `_msg` attribute. We can take a look at the first 4 and last 4 bytes of the packages message to see that it is a complete messages.  The first 4 bytes should be `'GRIB'` and last 4 `'7777'`.
# 
# **NOTE:** When we print the message, notice the message number of -1. This is because this Grib2Message object is not associated with a GRIB2 file object (`grib2io.open`).

# In[57]:


print(wspd._msg[0:4],wspd._msg[-4:])
print(wspd)
wspd


# ## **Writing Grib2Message Objects to a New File**
# 
# Now that we have created a new Grib2Message object, we can write that message to a new file.
# 
# Here we create a new GRIB2 file object to write to

# In[59]:


gout = grib2io.open('wind-speed-demo.grib2',mode='w')
print(gout)


# We can use the `write()` method to write a Grib2Message object to file.  We will use the wind speed message we created above.

# In[60]:


gout.write(wspd)


# In[61]:


print(gout)


# **NOTE:** Even though we wrote to the new file, the `variables` and `levels` attributes have not been updated.  This feature will be added in the near future.

# In[63]:


gout.close()
g.close()