#!/usr/bin/env python3 

# Make ensemble mean

import pywgrib2_s

in_file='c.grb'
out_file='c_ensmean.grb'
n_ens = 3

# Step 1. make sample ensemble data

file = {}
for i in range(n_ens):
    j = i + 1
    file[i] = "c_ensm%03d.grb" % j
    pywgrib2_s.wgrib2( [in_file, "-rewind_init", in_file, "-set_ens_num", "3", str(j), str(n_ens), 
      "-grib_out", file[i], "-inv", "/dev/null" ] )
print("list of files:", file)

# Step 2. make inventory of ensemble members

for i in range(n_ens):
    err=pywgrib2_s.mk_inv(file[i],file[i]+'.inv')
    print("mkinv file ",i," err=",err)

# Step 3. now to make ensemble mean, read file[0]+'.inv' sequentially

sequential=0
while pywgrib2_s.inq(file[0], inv=file[0]+'.inv', sequential=sequential, Data=True, Matched=True, grib='@mem:0') == 1:
    metadata = pywgrib2_s.matched[0]
    sum = pywgrib2_s.data

    # remove columns 1 and 2 to get prelim search string
    i = metadata.find(":", metadata.find(":")+1)+1
    search = metadata[i:]
    for i in range(1,n_ens):
        j = i + 1
        search_mem = search.replace('ENS=+1','ENS=+'+str(j))
        nmatch = pywgrib2_s.inq(file[i], search_mem, inv=file[i]+'.inv', Data=True)
        if nmatch != 1:
            print("problem with finding data=",search_mem)
            print("in file, ", file[i])
            quit()
        sum = sum + pywgrib2_s.data

    # write ensemble mean
    sum = sum / n_ens
    ens_mean_meta=search.replace('ENS=+1','ens mean')
    a=pywgrib2_s.write(out_file,'@mem:0',1,metadata=ens_mean_meta)
    if len(a) == 0:
        print("write error")
        quit()
    print("wrote ",ens_mean_meta)
    sequential=sequential+1

print("done count=",sequential, " output in ",out_file)