#!/usr/bin/env python3

# subset: random access
#
# This version is even faster than cookbook_subset_4.py and cookbook_subset_5.py
#
# this version creates a list of fields (new_inv) to copy from
# in_file to out.  Then calls pywgrib2_s.inq(..) to process the
# entire list.  This is faster than cookbook_subset_(4|5).py
# because there is only one call to wgrib2.


def find_in_list(string,inv):
    # return first string with match
    n=0
    j=0
    for i in range(len(inv)):
        if inv[i].find(string) != -1:
            n=n+1
            j=i
    if n == 1:
       return inv[j]
    if n == 0:
       print("No match!")
    if n > 1:
       print("Multple matches!")
    return None

import pywgrib2_s
in_file='a.grb'
out='a_subset_6.grb'
inv_file='@mem:0'

# make inventory file and save it in inv[]
# can use the short inventory it has all the info needed for problem

ierr = pywgrib2_s.mk_inv(in_file, inv_file, Short=True)
print('mk_inv ierr=',ierr)
inv = pywgrib2_s.read_inv(inv_file)
print("read inventory, messages=", len(inv))

# make an inventory of selected fields

new_inv=''
for field in (':TMP:2 m above ground:', ':PRATE:', ':UGRD:10 m above ground:'):
    for date in (20200101000000,20200101060000,20200101120000,20200101180000):
        field_date=':D='+str(date)+field
        print('processing: ',  field_date)
        record=find_in_list(field_date, inv)
        if record is None:
            print("Error did not find ", field_date)
            quit()
        new_inv=new_inv+record+'\n'

print("new_inv=")
print(new_inv)

# the inventory needs to be in a file to use
# save new_inv in memory file 0

ierr = pywgrib2_s.set_mem(0,new_inv)
print("write inv to @mem:0 error=",ierr)

# do an inquire with no search terms using the selected inventory
count = pywgrib2_s.inq(in_file,inv='@mem:0',grib=out)

print("done count = ", count, ' out=',out)