#!/usr/bin/env python3

# calculate the wind speed for u/v
# read U sequentially

import pywgrib2_s
import numpy

in_grib = "a.grb"
inv1='a.grb.inv1'
inv2='a.grb.inv2'
out='a_wind.grb'

# make two inventory files
err = pywgrib2_s.mk_inv(in_grib,inv1)
print("make inventory 1 err=", err)
err = pywgrib2_s.mk_inv(in_grib,inv2)
print("make inventory 2 err=", err)


count = 0
while pywgrib2_s.inq(in_grib,'UGRD',inv=inv1,Data=True, sequential=count, Matched=True) == 1:
    count += 1

    # save U
    u = pywgrib2_s.data

    # get match string for V
    metadata = pywgrib2_s.matched[0]

    # metadata UGRD -> VGRD, remove columns 1 and 2 (otherwise will not match)
    metadata = metadata.replace(":UGRD:",":VGRD:")
    i = metadata.find(":", metadata.find(":")+1)+1
    metadata = metadata[i:]

    # read corresponding V
    nmatch = pywgrib2_s.inq(in_grib, metadata, inv=inv2, Data=True, grib='@mem:0')
    if nmatch != 1:
        print("error looking for V nmatch=",nmatch)
        quit()
    v = pywgrib2_s.data
    wind = numpy.sqrt(u*u + v*v)

    # write write wind speed, modify the V field
    s=pywgrib2_s.write(out,'@mem:0',1,new_data=wind,var='WIND')
    if s is None:
        print("write failed")
        count -= 1


pywgrib2_s.close(out)

print("wrote ",count," wind fields to ",out)
print("\n  Note: if you get error message: RuntimeWarning: invalid value encountered in sqrt")
print("    You may be using a buggy anaconda-numpy.")
print("    See  https://github.com/numpy/numpy/issues/11448")
print("    fix: install original numpy")