#!/usr/bin/env python3

# compute the average RH from the surface to 700 mb and write as a grib file
#
# input: grib file with
#          surface pressure
#          RH on 1000, 925, 850 and 700 mb levels
#           (yes typical files has more levels of RH, but only example)
#
# output: grib file with pressure weighted RH from 700mb to surface
#
# for RH below 1000 mb, use 1000 mb value
#

import pywgrib2_s
import numpy

# can use regular files for inv and grib, but memory files are faster
#  and don't have be cleaned up afterwards

input="c.grb"
inv="@mem:0"
out='c_ave_rh.grb'
grib='@mem:1'
level = [ 1000, 925, 850, 700 ]

# make inventory file
pywgrib2_s.mk_inv(input, inv)
print("made inventory file")

# read surface pressure
n = pywgrib2_s.inq(input,":PRES:surface:",inv=inv,Data=True,grib=grib,Matched=True)
if n != 1:
    print("problem with reading surface pressure, number of matches=",n)
    quit()

# save pywgrib2_s.data before it get overwritten by next read
# convert surface pressure from Pa to mb (hPa)
ps = 0.01 * pywgrib2_s.data
print("surface pressure ",pywgrib2_s.matched[0])

# find shape of ps
if ps.ndim != 2:
    print("expected NxM array, ndim=",ps.ndim)
    quit()
nx=ps.shape[0]
ny=ps.shape[1]

# allocate numpy 3d array
rh = numpy.zeros((nx, ny, len(level)), order='F')

print("number of levels=",len(level))
print("shape rh=",rh.shape)
for i, lev in enumerate(level):
#   make lev compatible with wgrib2 search
    slev=":"+str(lev)+" mb:"
    n = pywgrib2_s.inq(input,":RH:",slev,inv=inv,Data=True)
    if n != 1: 
        print("problem with RH number of matches=",n," level=",level)
        quit()
    rh[:,:,i] = pywgrib2_s.data[:,:]
    print("read in RH lev=",lev)
print("read all RH fields")

print(ps > 1000.0)

# do the layer 1000mb - ps

layer_rh = rh[:,:,0]
ave_rh = numpy.maximum(ps-1000.0,0.0)*layer_rh
print('ave_rh')
print(ave_rh)

for layer in range(len(level)-1):
#   make lev compatible with wgrib2 search
    layer_rh = 0.5 * (rh[:,:,layer] + rh[:,:,layer+1])
    level0=level[layer]
    level1=level[layer+1]
    print("layer=", layer, level0,' to ',level1)
    ave_rh += numpy.minimum(numpy.maximum(ps-level1,0.0),level0-level1)*layer_rh

# divide by layer thickness, note: ave_rh is non-zero for thickness > 0

top = level[len(level)-1]
print("top=",top)
thickness = ps-top
print("thickness")
print(thickness)
thickness[thickness == 0.0] = 1.0
print("ave_rh")
ave_rh = ave_rh / thickness
print(ave_rh)