#!/usr/bin/env python3

# compute daily mean/max forecasts

import pywgrib2_s
import numpy as np

out='b_ave.grb'
nday=2

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

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

# pywgrib2_s.debug=True

# process ftime in format:  N-M hour ave fcst
print("process PRATE")
for var in [ 'PRATE' ] :
  for fday in range(1,3) :
    print("fday=",fday)
    for fhour in range( (fday-1)*24, fday*24, 6) :
      fhour2 = fhour+6
      ftime=str(fhour)+'-'+str(fhour2)+' hour ave fcst'
      i = pywgrib2_s.inq(file[fhour2],inv=file[fhour2]+'.inv',var=var,ftime=ftime,Matched=True,Data=True)
      print("#matches=",i, " read=",pywgrib2_s.matched[0])
      if fhour == (fday-1)*24 :
        tmp = pywgrib2_s.data
      else :
        tmp = tmp + pywgrib2_s.data
    tmp = tmp * 0.25
    ftime=str(fday*24-24)+'-'+str(fday*24)+' hour ave fcst'
    print("output ftime=",ftime)
    s = pywgrib2_s.write(out,file[0],1,new_data=tmp,metadata=pywgrib2_s.matched[0],ftime=ftime)
    print("wrote: ", s)

print("finished PRATE")

# process ftime in format:  N hour fcst
for var in [ 'UGRD', 'VGRD' ]:
  for fday in range(1,3) :
    print("fday=",fday)
    for fhour in range( (fday-1)*24, fday*24, 6) :
      ftime=str(fhour)+' hour fcst'
      if fhour == 0:
        ftime='anl'
      i = pywgrib2_s.inq(file[fhour],inv=file[fhour]+'.inv',var=var,ftime=ftime,Matched=True,Data=True)
      if i != 1:
        print("could not read file=",file[fhour]," ftime=",ftime)
        quit()
      print("#matches=",i, " read=",pywgrib2_s.matched[0])
      if fhour == (fday-1)*24 :
        tmp = pywgrib2_s.data
      else :
        tmp = tmp + pywgrib2_s.data
    tmp = tmp * 0.25
    ftime=str(fday*24-24)+'-'+str(fday*24-6)+' hour ave@(fcst,dt=6 hour)'
    s = pywgrib2_s.write(out,file[0],1,new_data=tmp,metadata=pywgrib2_s.matched[0],ftime=ftime)
    print("wrote: ", s)

print("finished U/V")

# process ftime in format:  N-M hour max fcst:
for var in [ 'TMAX' ]:
  for fday in range(1,3):
    print("fday=",fday)
    for fhour in range( (fday-1)*24, fday*24, 6) :
      fhour2 = fhour+6
      ftime=str(fhour)+'-'+str(fhour2)+' hour max fcst'
      i = pywgrib2_s.inq(file[fhour2],inv=file[fhour2]+'.inv',var=var,ftime=ftime,Matched=True,Data=True)
      if i != 1:
          print("problem reading file=",file[fhour2]," ftime=",ftime)
      print("#matches=",i, " read=",pywgrib2_s.matched[0])
      if fhour == (fday-1)*24 :
        tmp = pywgrib2_s.data
      else :
        tmp = np.maximum(tmp,pywgrib2_s.data)
    ftime=str(fday*24-24)+'-'+str(fday*24)+' hour max fcst'
    s = pywgrib2_s.write(out,file[0],1,new_data=tmp,metadata=pywgrib2_s.matched[0],ftime=ftime)
    print("wrote: ", s)
print("finished TMAX")


print("done output in",out)