#!/usr/bin/env python3

# have forecast files with 6 hour spacing (6n hours)
# make time series with 3 hour spacing
# use linear interpolation for 6n+3 hours values
# variables must have ftime type 'N hour fcst'

def new_ftime(fhr):
# return ftime for given forecast hour
    if fhr == 0:
        return 'anl'
    else:
        return str(fhr)+' hour fcst' 

import pywgrib2_s
import numpy as np

out='b_3hr.grb'
time_start = 0
time_end = 48
vars = [ 'UGRD:10 m above ground', 'VGRD:10 m above ground' ]


# make dictionary of input file names
file = {}
for fhr in range(time_start, time_end+1,6):
    ff = "b.f" "%03d" % fhr
    file[fhr] = ff
print("list of files:", file)

# make inventory files
for f in range(time_start,time_end+1,6):
    pywgrib2_s.mk_inv(file[f], file[f]+'.inv')
print("finished making inv files")


# for fhr = 6n hours, copy grib message
# for fhr = 6n + 3 hours, interpolate 6n and 6n+6 and write out as 6n+3 hours

for fhr in range(time_start, time_end+1, 3):
    if fhr % 6 == 0:   # fhr=0,6,12,18,...
        for var in vars:
        # copy message to output
            i = pywgrib2_s.inq(file[fhr],var=var,inv=file[fhr]+'.inv',grib=out,Matched=True,ftime=new_ftime(fhr))
            print("copied: #match=",i,pywgrib2_s.matched[0])
    else:             # fhr=3,9,15..
        for var in vars :
            # compute 6n+3 by interpolating 6n and 6n+6
            i = pywgrib2_s.inq(file[fhr-3],var=var,inv=file[fhr-3]+'.inv',Data=True)
            tmp = pywgrib2_s.data
            # save grib message as @mem:0, use as template
            j = pywgrib2_s.inq(file[fhr+3],var=var,inv=file[fhr+3]+'.inv',Data=True,Matched=True,grib='@mem:0')
            tmp = (tmp+pywgrib2_s.data)*0.5
            s = pywgrib2_s.write(out,'@mem:0',1,new_data=tmp,ftime=new_ftime(fhr))
            if s is None:
                print("write failed")
                quit()
            print("computed  ",s)

print("done ",out)